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.
  """Base class for entrypoint configurations that run a single step."""
  
  import os
  import sys
  from typing import TYPE_CHECKING, Any, Dict, List
  from uuid import UUID
  
  from zenml.client import Client
  from zenml.entrypoints.base_entrypoint_configuration import (
      SNAPSHOT_ID_OPTION,
      BaseEntrypointConfiguration,
  )
  from zenml.integrations.registry import integration_registry
  from zenml.logger import get_logger
  
  if TYPE_CHECKING:
      from zenml.config import DockerSettings
      from zenml.config.step_configurations import Step
      from zenml.models import PipelineSnapshotResponse
  
  logger = get_logger(__name__)
  
  STEP_NAME_OPTION = "step_name"
  
  
  class StepEntrypointConfiguration(BaseEntrypointConfiguration):
      """Base class for entrypoint configurations that run a single step.
  
      If an orchestrator needs to run steps in a separate process or environment
      (e.g. a docker container), this class can either be used directly or
      subclassed if custom behavior is necessary.
  
      How to subclass:
      ----------------
      Passing additional arguments to the entrypoint:
          If you need to pass additional arguments to the entrypoint, there are
          two methods that you need to implement:
              * `get_entrypoint_options()`: This method should return all
                  the options that are required in the entrypoint. Make sure to
                  include the result from the superclass method so the options
                  are complete.
  
              * `get_entrypoint_arguments(...)`: This method should return
                  a list of arguments that should be passed to the entrypoint.
                  Make sure to include the result from the superclass method so
                  the arguments are complete.
  
          You'll be able to access the argument values from `self.entrypoint_args`
          inside your `StepEntrypointConfiguration` subclass.
  
      How to use:
      -----------
      After you created your `StepEntrypointConfiguration` subclass, you only
      have to run the entrypoint somewhere. To do this, you should execute the
      command returned by the `get_entrypoint_command()` method with the
      arguments returned by the `get_entrypoint_arguments(...)` method.
  
      Example:
      ```python
      class MyStepEntrypointConfiguration(StepEntrypointConfiguration):
          ...
  
      class MyOrchestrator(BaseOrchestrator):
          def submit_pipeline(
              self,
              snapshot: "PipelineSnapshotResponse",
              stack: "Stack",
              environment: Dict[str, str],
              placeholder_run: Optional["PipelineRunResponse"] = None,
          ) -> Optional[SubmissionResult]:
              ...
  
              cmd = MyStepEntrypointConfiguration.get_entrypoint_command()
              for step_name, step in pipeline.steps.items():
                  ...
  
                  args = MyStepEntrypointConfiguration.get_entrypoint_arguments(
                      step_name=step_name, snapshot_id=snapshot.id
                  )
                  # Run the command and pass it the arguments. Our example
                  # orchestrator here executes the entrypoint in a separate
                  # process, but in a real-world scenario you would probably run
                  # it inside a docker container or a different environment.
                  import subprocess
                  subprocess.check_call(cmd + args)
      ```
      """
  
      def post_run(
          self,
          pipeline_name: str,
          step_name: str,
      ) -> None:
          """Does cleanup or post-processing after the step finished running.
  
          Subclasses should overwrite this method if they need to run any
          additional code after the step execution.
  
          Args:
              pipeline_name: Name of the parent pipeline of the step that was
                  executed.
              step_name: Name of the step that was executed.
          """
  
      @classmethod
      def get_entrypoint_options(cls) -> Dict[str, bool]:
          """Gets all options required for running with this configuration.
  
          Returns:
              The superclass options as well as an option for the name of the
              step to run.
          """
          return super().get_entrypoint_options() | {STEP_NAME_OPTION: True}
  
      @classmethod
      def get_entrypoint_arguments(
          cls,
          **kwargs: Any,
      ) -> List[str]:
          """Gets all arguments that the entrypoint command should be called with.
  
          The argument list should be something that
          `argparse.ArgumentParser.parse_args(...)` can handle (e.g.
          `["--some_option", "some_value"]` or `["--some_option=some_value"]`).
          It needs to provide values for all options returned by the
          `get_entrypoint_options()` method of this class.
  
          Args:
              **kwargs: Kwargs, must include the step name.
  
          Returns:
              The superclass arguments as well as arguments for the name of the
              step to run.
          """
          return super().get_entrypoint_arguments(**kwargs) + [
              f"--{STEP_NAME_OPTION}",
              kwargs[STEP_NAME_OPTION],
          ]
  
      @property
      def docker_settings(self) -> "DockerSettings":
          """The Docker settings configured for this entrypoint configuration.
  
          Returns:
              The Docker settings.
          """
          return self.step.config.docker_settings
  
      @property
      def step(self) -> "Step":
          """The step configured for this entrypoint configuration.
  
          Returns:
              The step.
          """
          step_name = self.entrypoint_args[STEP_NAME_OPTION]
          return self.snapshot.step_configurations[step_name]
  
      def _load_snapshot(self) -> "PipelineSnapshotResponse":
          """Loads the snapshot.
  
          Returns:
              The snapshot.
          """
          snapshot_id = UUID(self.entrypoint_args[SNAPSHOT_ID_OPTION])
          step_name = self.entrypoint_args[STEP_NAME_OPTION]
          return Client().zen_store.get_snapshot(
              snapshot_id=snapshot_id, step_configuration_filter=[step_name]
          )
  
      def run(self) -> None:
          """Prepares the environment and runs the configured step."""
          snapshot = self.snapshot
  
          # Activate all the integrations. This makes sure that all materializers
          # and stack component flavors are registered.
          integration_registry.activate_integrations()
  
          step_name = self.entrypoint_args[STEP_NAME_OPTION]
  
          # Change the working directory to make sure we're in the correct
          # directory where the files in the Docker image should be included.
          # This is necessary as some services overwrite the working directory
          # configured in the Docker image itself.
          os.makedirs("/app", exist_ok=True)
          os.chdir("/app")
  
          self.download_code_if_necessary()
  
          # If the working directory is not in the sys.path, we include it to make
          # sure user code gets correctly imported.
          cwd = os.getcwd()
          if cwd not in sys.path:
              sys.path.insert(0, cwd)
  
          pipeline_name = snapshot.pipeline_configuration.name
  
          self._run_step(step=self.step, snapshot=snapshot)
  
          self.post_run(
              pipeline_name=pipeline_name,
              step_name=step_name,
          )
  
      def _run_step(
          self,
          step: "Step",
          snapshot: "PipelineSnapshotResponse",
      ) -> None:
          """Runs a single step.
  
          Args:
              step: The step to run.
              snapshot: The snapshot configuration.
          """
          orchestrator = Client().active_stack.orchestrator
          orchestrator._prepare_run(snapshot=snapshot)
          orchestrator.run_step(step=step)
  