content: #  Copyright (c) ZenML GmbH 2022. 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:
  #
  #       http://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.
  """Step context class."""
  
  import contextvars
  from typing import (
      TYPE_CHECKING,
      Any,
      Dict,
      List,
      Mapping,
      Optional,
      Sequence,
      Type,
  )
  
  from typing_extensions import Self
  
  from zenml.exceptions import StepContextError
  from zenml.logger import get_logger
  from zenml.utils import context_utils
  from zenml.utils.callback_registry import CallbackRegistry
  from zenml.utils.singleton import SingletonMetaClass
  
  if TYPE_CHECKING:
      from zenml.artifacts.artifact_config import ArtifactConfig
      from zenml.materializers.base_materializer import BaseMaterializer
      from zenml.metadata.metadata_types import MetadataType
      from zenml.model.model import Model
      from zenml.models import (
          PipelineResponse,
          PipelineRunResponse,
          StepRunResponse,
      )
      from zenml.models.v2.core.step_run import StepRunInputResponse
  
  
  logger = get_logger(__name__)
  
  
  def get_step_context() -> "StepContext":
      """Get the context of the currently running step.
  
      Returns:
          The context of the currently running step.
  
      Raises:
          RuntimeError: If no step is currently running.
      """
      if ctx := StepContext.get():
          return ctx
  
      raise RuntimeError(
          "The step context is only available inside a step function."
      )
  
  
  def get_or_create_run_context() -> "RunContext":
      """Get or create the context of the currently running pipeline.
  
      Returns:
          The context of the currently running pipeline.
      """
      return RunContext()
  
  
  class RunContext(metaclass=SingletonMetaClass):
      """Provides context shared between all steps in a pipeline run."""
  
      def __init__(self) -> None:
          """Create the run context."""
          self.initialized = False
          self._state: Optional[Any] = None
  
      @property
      def state(self) -> Optional[Any]:
          """Returns the pipeline state.
  
          Returns:
              The pipeline state or None.
  
          Raises:
              RuntimeError: If the run context is not initialized.
          """
          if not self.initialized:
              raise RuntimeError(
                  "Run context not initialized. The run state is only available "
                  "in the context of a running pipeline."
              )
          return self._state
  
      def initialize(self, state: Optional[Any]) -> None:
          """Initialize the run context.
  
          Args:
              state: Optional state for the pipeline run
  
          Raises:
              RuntimeError: If the run context is already initialized.
          """
          if self.initialized:
              raise RuntimeError("Run context already initialized.")
          self._state = state
          self.initialized = True
  
  
  class StepContext(context_utils.BaseContext):
      """Provides additional context inside a step function.
  
      Usage example:
  
      ```python
      from zenml.steps import get_step_context
  
      @step
      def my_trainer_step() -> Any:
          context = get_step_context()
  
          # get info about the current pipeline run
          current_pipeline_run = context.pipeline_run
  
          # get info about the current step run
          current_step_run = context.step_run
  
          # get info about the future output artifacts of this step
          output_artifact_uri = context.get_output_artifact_uri()
  
          ...
      ```
      """
  
      __context_var__ = contextvars.ContextVar("step_context")
  
      def __init__(
          self,
          pipeline_run: "PipelineRunResponse",
          step_run: "StepRunResponse",
          output_materializers: Mapping[str, Sequence[Type["BaseMaterializer"]]],
          output_artifact_uris: Mapping[str, str],
          output_artifact_configs: Mapping[str, Optional["ArtifactConfig"]],
      ) -> None:
          """Initialize the context of the currently running step.
  
          Args:
              pipeline_run: The model of the current pipeline run.
              step_run: The model of the current step run.
              output_materializers: The output materializers of the step that
                  this context is used in.
              output_artifact_uris: The output artifacts of the step that this
                  context is used in.
              output_artifact_configs: The outputs' ArtifactConfigs of the step that this
                  context is used in.
  
          Raises:
              StepContextError: If the keys of the output materializers and
                  output artifacts do not match.
          """
          from zenml.client import Client
  
          super().__init__()
  
          try:
              pipeline_run = Client().get_pipeline_run(pipeline_run.id)
          except KeyError:
              pass
          self.pipeline_run = pipeline_run
          try:
              step_run = Client().get_run_step(step_run.id)
          except KeyError:
              pass
          self.step_run = step_run
          self.model_version = (
              step_run.model_version or pipeline_run.model_version
          )
  
          self.step_name = self.step_run.name
  
          # set outputs
          if output_materializers.keys() != output_artifact_uris.keys():
              raise StepContextError(
                  f"Mismatched keys in output materializers and output artifact "
                  f"URIs for step `{self.step_name}`. Output materializer "
                  f"keys: {set(output_materializers)}, output artifact URI "
                  f"keys: {set(output_artifact_uris)}"
              )
          self._outputs = {
              key: StepContextOutput(
                  materializer_classes=output_materializers[key],
                  artifact_uri=output_artifact_uris[key],
                  artifact_config=output_artifact_configs[key],
              )
              for key in output_materializers.keys()
          }
          self._cleanup_registry = CallbackRegistry()
  
      @property
      def pipeline(self) -> "PipelineResponse":
          """Returns the current pipeline.
  
          Returns:
              The current pipeline or None.
  
          Raises:
              StepContextError: If the pipeline run does not have a pipeline.
          """
          if self.pipeline_run.pipeline:
              return self.pipeline_run.pipeline
          raise StepContextError(
              f"Unable to get pipeline in step `{self.step_name}` of pipeline "
              f"run '{self.pipeline_run.id}': This pipeline run does not have "
              f"a pipeline associated with it."
          )
  
      @property
      def pipeline_state(self) -> Optional[Any]:
          """Returns the pipeline state.
  
          Returns:
              The pipeline state or None.
          """
          return get_or_create_run_context().state
  
      @property
      def model(self) -> "Model":
          """Returns configured Model.
  
          Order of resolution to search for Model is:
              1. Model from the step context
              2. Model from the pipeline context
  
          Returns:
              The `Model` object associated with the current step.
  
          Raises:
              StepContextError: If no `Model` object was specified for the step
                  or pipeline.
          """
          if not self.model_version:
              raise StepContextError(
                  f"Unable to get Model in step `{self.step_name}` of pipeline "
                  f"run '{self.pipeline_run.id}': No model has been specified "
                  "the step or pipeline."
              )
  
          return self.model_version.to_model_class()
  
      @property
      def inputs(self) -> Dict[str, List["StepRunInputResponse"]]:
          """Returns the input artifacts of the current step.
  
          Returns:
              The input artifacts of the current step.
          """
          return self.step_run.regular_inputs
  
      def _get_output(
          self, output_name: Optional[str] = None
      ) -> "StepContextOutput":
          """Returns the materializer and artifact URI for a given step output.
  
          Args:
              output_name: Optional name of the output for which to get the
                  materializer and URI.
  
          Returns:
              Tuple containing the materializer and artifact URI for the
                  given output.
  
          Raises:
              StepContextError: If the step has no outputs, no output for
                  the given `output_name` or if no `output_name` was given but
                  the step has multiple outputs.
          """
          output_count = len(self._outputs)
          if output_count == 0:
              raise StepContextError(
                  f"Unable to get step output for step `{self.step_name}`: "
                  f"This step does not have any outputs."
              )
  
          if not output_name and output_count > 1:
              raise StepContextError(
                  f"Unable to get step output for step `{self.step_name}`: "
                  f"This step has multiple outputs ({set(self._outputs)}), "
                  f"please specify which output to return."
              )
  
          if output_name:
              if output_name not in self._outputs:
                  raise StepContextError(
                      f"Unable to get step output '{output_name}' for "
                      f"step `{self.step_name}`. This step does not have an "
                      f"output with the given name, please specify one of the "
                      f"available outputs: {set(self._outputs)}."
                  )
              return self._outputs[output_name]
          else:
              return next(iter(self._outputs.values()))
  
      def get_output_materializer(
          self,
          output_name: Optional[str] = None,
          custom_materializer_class: Optional[Type["BaseMaterializer"]] = None,
          data_type: Optional[Type[Any]] = None,
      ) -> "BaseMaterializer":
          """Returns a materializer for a given step output.
  
          Args:
              output_name: Optional name of the output for which to get the
                  materializer. If no name is given and the step only has a
                  single output, the materializer of this output will be
                  returned. If the step has multiple outputs, an exception
                  will be raised.
              custom_materializer_class: If given, this `BaseMaterializer`
                  subclass will be initialized with the output artifact instead
                  of the materializer that was registered for this step output.
              data_type: If the output annotation is of type `Union` and the step
                  therefore has multiple materializers configured, you can provide
                  a data type for the output which will be used to select the
                  correct materializer. If not provided, the first materializer
                  will be used.
  
          Returns:
              A materializer initialized with the output artifact for
              the given output.
          """
          from zenml.utils import materializer_utils
  
          output = self._get_output(output_name)
          materializer_classes = output.materializer_classes
          artifact_uri = output.artifact_uri
  
          if custom_materializer_class:
              materializer_class = custom_materializer_class
          elif len(materializer_classes) == 1 or not data_type:
              materializer_class = materializer_classes[0]
          else:
              materializer_class = materializer_utils.select_materializer(
                  data_type=data_type, materializer_classes=materializer_classes
              )
  
          return materializer_class(artifact_uri)
  
      def get_output_artifact_uri(
          self, output_name: Optional[str] = None
      ) -> str:
          """Returns the artifact URI for a given step output.
  
          Args:
              output_name: Optional name of the output for which to get the URI.
                  If no name is given and the step only has a single output,
                  the URI of this output will be returned. If the step has
                  multiple outputs, an exception will be raised.
  
          Returns:
              Artifact URI for the given output.
          """
          return self._get_output(output_name).artifact_uri
  
      def get_output_metadata(
          self, output_name: Optional[str] = None
      ) -> Dict[str, "MetadataType"]:
          """Returns the metadata for a given step output.
  
          Args:
              output_name: Optional name of the output for which to get the
                  metadata. If no name is given and the step only has a single
                  output, the metadata of this output will be returned. If the
                  step has multiple outputs, an exception will be raised.
  
          Returns:
              Metadata for the given output.
          """
          output = self._get_output(output_name)
          custom_metadata = output.run_metadata or {}
          if output.artifact_config:
              custom_metadata.update(
                  **(output.artifact_config.run_metadata or {})
              )
          return custom_metadata
  
      def get_output_tags(self, output_name: Optional[str] = None) -> List[str]:
          """Returns the tags for a given step output.
  
          Args:
              output_name: Optional name of the output for which to get the
                  metadata. If no name is given and the step only has a single
                  output, the metadata of this output will be returned. If the
                  step has multiple outputs, an exception will be raised.
  
          Returns:
              Tags for the given output.
          """
          output = self._get_output(output_name)
          custom_tags = set(output.tags or [])
          if output.artifact_config:
              return list(
                  set(output.artifact_config.tags or []).union(custom_tags)
              )
          return list(custom_tags)
  
      def add_output_metadata(
          self,
          metadata: Dict[str, "MetadataType"],
          output_name: Optional[str] = None,
      ) -> None:
          """Adds metadata for a given step output.
  
          Args:
              metadata: The metadata to add.
              output_name: Optional name of the output for which to add the
                  metadata. If no name is given and the step only has a single
                  output, the metadata of this output will be added. If the
                  step has multiple outputs, an exception will be raised.
          """
          output = self._get_output(output_name)
          if not output.run_metadata:
              output.run_metadata = {}
          output.run_metadata.update(**metadata)
  
      def add_output_tags(
          self,
          tags: List[str],
          output_name: Optional[str] = None,
      ) -> None:
          """Adds tags for a given step output.
  
          Args:
              tags: The tags to add.
              output_name: Optional name of the output for which to add the
                  tags. If no name is given and the step only has a single
                  output, the tags of this output will be added. If the
                  step has multiple outputs, an exception will be raised.
          """
          output = self._get_output(output_name)
          if not output.tags:
              output.tags = []
          output.tags += tags
  
      def remove_output_tags(
          self,
          tags: List[str],
          output_name: Optional[str] = None,
      ) -> None:
          """Removes tags for a given step output.
  
          Args:
              tags: The tags to remove.
              output_name: Optional name of the output for which to remove the
                  tags. If no name is given and the step only has a single
                  output, the tags of this output will be removed. If the
                  step has multiple outputs, an exception will be raised.
          """
          output = self._get_output(output_name)
          if not output.tags:
              return
          output.tags = [tag for tag in output.tags if tag not in tags]
  
      def __enter__(self) -> Self:
          """Enter the step context.
  
          Raises:
              RuntimeError: If the step context has already been entered.
  
          Returns:
              The step context object.
          """
          if self._token is not None:
              raise RuntimeError(
                  "Running a step from within another step is not allowed."
              )
          return super().__enter__()
  
      def __exit__(self, *_: Any) -> None:
          """Exit the step context.
  
          Args:
              *_: Unused keyword arguments.
          """
          self._cleanup_registry.execute_callbacks(raise_on_exception=False)
          super().__exit__(*_)
  
  
  class StepContextOutput:
      """Represents a step output in the step context."""
  
      materializer_classes: Sequence[Type["BaseMaterializer"]]
      artifact_uri: str
      run_metadata: Optional[Dict[str, "MetadataType"]] = None
      artifact_config: Optional["ArtifactConfig"]
      tags: Optional[List[str]] = None
  
      def __init__(
          self,
          materializer_classes: Sequence[Type["BaseMaterializer"]],
          artifact_uri: str,
          artifact_config: Optional["ArtifactConfig"],
      ):
          """Initialize the step output.
  
          Args:
              materializer_classes: The materializer classes for the output.
              artifact_uri: The artifact URI for the output.
              artifact_config: The ArtifactConfig object of the output.
          """
          self.materializer_classes = materializer_classes
          self.artifact_uri = artifact_uri
          self.artifact_config = artifact_config
  