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:
  #
  #       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.
  """Utilities for caching."""
  
  import hashlib
  import os
  from typing import TYPE_CHECKING, List, Mapping, Optional
  from uuid import UUID
  
  from zenml.client import Client
  from zenml.constants import CODE_HASH_PARAMETER_NAME
  from zenml.enums import ExecutionStatus, SorterOps
  from zenml.logger import get_logger
  from zenml.orchestrators import step_run_utils
  from zenml.utils import source_code_utils, source_utils
  
  if TYPE_CHECKING:
      from zenml.artifact_stores import BaseArtifactStore
      from zenml.config.step_configurations import Step
      from zenml.models import (
          PipelineRunResponse,
          PipelineSnapshotResponse,
          StepRunResponse,
      )
      from zenml.models.v2.core.step_run import (
          StepRunInputResponse,
      )
      from zenml.stack import Stack
  
  
  logger = get_logger(__name__)
  
  
  def generate_cache_key(
      step: "Step",
      input_artifacts: Mapping[str, List["StepRunInputResponse"]],
      artifact_store: "BaseArtifactStore",
      project_id: "UUID",
  ) -> str:
      """Generates a cache key for a step run.
  
      If a step has caching enabled, the cache key will be used to find existing
      equivalent step runs.
  
      The cache key is a MD5 hash of:
      - the project ID,
      - the artifact store ID and path,
      - the source code that defines the step,
      - the parameters of the step,
      - the names and IDs of the input artifacts of the step,
      - the names and source codes of the output artifacts of the step,
      - the source codes of the output materializers of the step.
      - additional custom caching parameters of the step.
      - the environment variables defined for the step.
      - the secrets defined for the step.
      - other elements of the cache policy defined for the step.
  
      Args:
          step: The step to generate the cache key for.
          input_artifacts: The input artifacts for the step.
          artifact_store: The artifact store of the active stack.
          project_id: The ID of the active project.
  
      Raises:
          ValueError: If some file dependencies are outside the source root or
              missing.
          ValueError: If the cache function is invalid.
          RuntimeError: If executing the cache function failed.
  
      Returns:
          A cache key.
      """
      cache_policy = step.config.cache_policy
      hash_ = hashlib.md5()  # nosec
  
      # Project ID
      hash_.update(project_id.bytes)
  
      # Artifact store ID and path
      hash_.update(artifact_store.id.bytes)
      hash_.update(artifact_store.path.encode())
  
      if artifact_store.custom_cache_key:
          hash_.update(artifact_store.custom_cache_key)
  
      # Step source. This currently only uses the string representation of the
      # source (e.g. my_module.step_class) instead of the full source to keep
      # the caching behavior of previous versions and to not invalidate caching
      # when committing some unrelated files
      hash_.update(step.spec.source.import_path.encode())
  
      if cache_policy.include_step_parameters:
          for key, value in sorted(step.config.parameters.items()):
              hash_.update(key.encode())
              hash_.update(str(value).encode())
  
      # Input artifacts
      for name, artifact_versions in input_artifacts.items():
          if name in (cache_policy.ignored_inputs or []):
              continue
  
          hash_.update(name.encode())
  
          for artifact_version in artifact_versions:
              if (
                  artifact_version.content_hash
                  and cache_policy.include_artifact_values
              ):
                  hash_.update(artifact_version.content_hash.encode())
              elif cache_policy.include_artifact_ids:
                  hash_.update(artifact_version.id.bytes)
  
              if artifact_version.chunk_index is not None:
                  hash_.update(str(artifact_version.chunk_index).encode())
                  chunk_size = artifact_version.chunk_size or 1
                  hash_.update(str(chunk_size).encode())
  
      # Output artifacts and materializers
      for name, output in step.config.outputs.items():
          hash_.update(name.encode())
          for source in output.materializer_source:
              hash_.update(source.import_path.encode())
  
      # Custom caching parameters
      for key, value in sorted(step.config.caching_parameters.items()):
          if (
              key == CODE_HASH_PARAMETER_NAME
              and not cache_policy.include_step_code
          ):
              continue
  
          hash_.update(key.encode())
          hash_.update(str(value).encode())
  
      # User-defined environment variables
      for key, value in sorted(step.config.environment.items()):
          hash_.update(key.encode())
          hash_.update(str(value).encode())
  
      # User-defined secrets
      for secret_name_or_id in sorted([str(s) for s in step.config.secrets]):
          hash_.update(secret_name_or_id.encode())
  
      if file_dependencies := cache_policy.file_dependencies:
          source_root = source_utils.get_source_root()
  
          absolute_paths = []
          for file_path in file_dependencies:
              if os.path.isabs(file_path):
                  if not os.path.relpath(file_path, source_root):
                      raise ValueError(
                          f"Cache policy file dependency `{file_path}` is not "
                          f"within your source root `{source_root}`. Only files "
                          "within your source root are allowed."
                      )
              else:
                  file_path = os.path.abspath(
                      os.path.join(source_root, file_path)
                  )
  
              if not os.path.exists(file_path):
                  raise ValueError(
                      f"Missing file for cache computation: `{file_path}`. "
                      "Please make sure to specify values in "
                      "`CachePolicy.file_dependencies` as a path relative to "
                      f"your source root ({source_root})."
                  )
  
              absolute_paths.append(file_path)
  
          for file_path in sorted(absolute_paths):
              # Use relative path to source root as key so we don't generate
              # different hashes on different machines
              relative_path = os.path.relpath(file_path, source_root)
              hash_.update(relative_path.encode())
              with open(file_path, "rb") as f:
                  while chunk := f.read(1024 * 1024):
                      hash_.update(chunk)
  
      if cache_policy.source_dependencies:
          source_dependencies = [
              source.import_path for source in cache_policy.source_dependencies
          ]
          for source_path in sorted(source_dependencies):
              hash_.update(source_path.encode())
              object_ = source_utils.load(source_path)
              source_code = source_code_utils.get_source_code(object_)
              hash_.update(source_code.encode())
  
      if cache_policy.cache_func:
          cache_func = source_utils.load(cache_policy.cache_func)
  
          if not callable(cache_func):
              raise ValueError(
                  f"The cache function {cache_policy.cache_func.import_path} is "
                  "not callable."
              )
  
          try:
              result = cache_func()
          except Exception as e:
              raise RuntimeError("Failed to run cache function.") from e
  
          if not isinstance(result, str):
              raise ValueError(
                  f"The cache function {cache_policy.cache_func.import_path} "
                  "must return a string."
              )
  
          hash_.update(result.encode())
  
      return hash_.hexdigest()
  
  
  def get_cached_step_run(cache_key: str) -> Optional["StepRunResponse"]:
      """If a given step can be cached, get the corresponding existing step run.
  
      A step run can be cached if there is an existing step run in the same
      project which has the same cache key and was successfully executed.
  
      Args:
          cache_key: The cache key of the step.
  
      Returns:
          The existing step run if the step can be cached, otherwise None.
      """
      client = Client()
  
      cache_candidates = client.list_run_steps(
          project=client.active_project.id,
          cache_key=cache_key,
          cache_expired=False,
          status=ExecutionStatus.COMPLETED,
          sort_by=f"{SorterOps.DESCENDING}:created",
          size=1,
      ).items
  
      if cache_candidates:
          return cache_candidates[0]
      return None
  
  
  def create_cached_step_runs_and_prune_snapshot(
      snapshot: "PipelineSnapshotResponse",
      pipeline_run: "PipelineRunResponse",
      stack: "Stack",
  ) -> bool:
      """Create cached step runs and prune the cached steps from the snapshot.
  
      Args:
          snapshot: The pipeline snapshot.
          pipeline_run: The pipeline run for which to create the step runs.
          stack: The stack on which the pipeline run is happening.
  
      Returns:
          Whether an actual pipeline run is still required.
      """
      cached_invocations = step_run_utils.create_cached_step_runs(
          snapshot=snapshot,
          pipeline_run=pipeline_run,
          stack=stack,
      )
  
      for invocation_id in cached_invocations:
          # Remove the cached step invocations from the snapshot so
          # the orchestrator does not try to run them
          snapshot.step_configurations.pop(invocation_id)
  
      for step in snapshot.step_configurations.values():
          for invocation_id in cached_invocations:
              if invocation_id in step.spec.upstream_steps:
                  step.spec.upstream_steps.remove(invocation_id)
  
      if len(snapshot.step_configurations) == 0:
          # All steps were cached, we update the pipeline run status and
          # don't actually use the orchestrator to run the pipeline
          logger.info("All steps of the pipeline run were cached.")
          return False
  
      return True
  