content: #  Copyright (c) ZenML GmbH 2023. 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.
  """Implementation of the Lightning orchestrator."""
  
  import os
  import tempfile
  import time
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast
  from uuid import uuid4
  
  from lightning_sdk import Machine, Studio
  
  from zenml import __version__
  from zenml.constants import (
      ENV_ZENML_CUSTOM_SOURCE_ROOT,
  )
  from zenml.entrypoints.step_entrypoint_configuration import (
      StepEntrypointConfiguration,
  )
  from zenml.integrations.lightning.flavors.lightning_orchestrator_flavor import (
      LightningOrchestratorConfig,
      LightningOrchestratorSettings,
  )
  from zenml.integrations.lightning.orchestrators.lightning_orchestrator_entrypoint_configuration import (
      LightningOrchestratorEntrypointConfiguration,
  )
  from zenml.integrations.lightning.orchestrators.utils import (
      gather_requirements,
      sanitize_studio_name,
  )
  from zenml.logger import get_logger
  from zenml.orchestrators import BaseOrchestrator, SubmissionResult
  from zenml.orchestrators.utils import get_orchestrator_run_name
  from zenml.utils import code_utils, io_utils, source_utils
  
  if TYPE_CHECKING:
      from zenml.models import PipelineRunResponse, PipelineSnapshotResponse
      from zenml.stack import Stack
  
  
  logger = get_logger(__name__)
  ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID = "ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID"
  LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH = "."
  
  
  class LightningOrchestrator(BaseOrchestrator):
      """Lightning orchestrator."""
  
      def _set_lightning_env_vars(
          self,
          snapshot: "PipelineSnapshotResponse",
      ) -> None:
          """Set up the Lightning client using environment variables.
  
          Args:
              snapshot: The pipeline snapshot to prepare or run.
  
          Raises:
              ValueError: If the user id and api key or username and organization
          """
          settings = cast(
              LightningOrchestratorSettings, self.get_settings(snapshot)
          )
          if not settings.user_id or not settings.api_key:
              raise ValueError(
                  "Lightning orchestrator requires `user_id` and `api_key` both to be set in the settings."
              )
          os.environ["LIGHTNING_USER_ID"] = settings.user_id
          os.environ["LIGHTNING_API_KEY"] = settings.api_key
          if settings.username:
              os.environ["LIGHTNING_USERNAME"] = settings.username
          elif settings.organization:
              os.environ["LIGHTNING_ORG"] = settings.organization
          else:
              raise ValueError(
                  "Lightning orchestrator requires either `username` or `organization` to be set in the settings."
              )
          if settings.teamspace:
              os.environ["LIGHTNING_TEAMSPACE"] = settings.teamspace
  
      @property
      def config(self) -> LightningOrchestratorConfig:
          """Returns the `LightningOrchestratorConfig` config.
  
          Returns:
              The configuration.
          """
          return cast(LightningOrchestratorConfig, self._config)
  
      @property
      def settings_class(self) -> Type[LightningOrchestratorSettings]:
          """Settings class for the Lightning orchestrator.
  
          Returns:
              The settings class.
          """
          return LightningOrchestratorSettings
  
      def get_orchestrator_run_id(self) -> str:
          """Returns the active orchestrator run id.
  
          Raises:
              RuntimeError: If no run id exists. This happens when this method
                  gets called while the orchestrator is not running a pipeline.
  
          Returns:
              The orchestrator run id.
  
          Raises:
              RuntimeError: If the run id cannot be read from the environment.
          """
          try:
              return os.environ[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID]
          except KeyError:
              raise RuntimeError(
                  "Unable to read run id from environment variable "
                  f"{ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID}."
              )
  
      @property
      def root_directory(self) -> str:
          """Path to the root directory for all files concerning this orchestrator.
  
          Returns:
              Path to the root directory.
          """
          return os.path.join(
              io_utils.get_global_config_directory(),
              "lightning",
              str(self.id),
          )
  
      @property
      def pipeline_directory(self) -> str:
          """Returns path to a directory in which the kubeflow pipeline files are stored.
  
          Returns:
              Path to the pipeline directory.
          """
          return os.path.join(self.root_directory, "pipelines")
  
      def setup_credentials(self) -> None:
          """Set up credentials for the orchestrator."""
          connector = self.get_connector()
          assert connector is not None
          connector.configure_local_client()
  
      def submit_pipeline(
          self,
          snapshot: "PipelineSnapshotResponse",
          stack: "Stack",
          base_environment: Dict[str, str],
          step_environments: Dict[str, Dict[str, str]],
          placeholder_run: Optional["PipelineRunResponse"] = None,
      ) -> Optional[SubmissionResult]:
          """Submits a pipeline to the orchestrator.
  
          This method should only submit the pipeline and not wait for it to
          complete. If the orchestrator is configured to wait for the pipeline run
          to complete, a function that waits for the pipeline run to complete can
          be passed as part of the submission result.
  
          Args:
              snapshot: The pipeline snapshot to submit.
              stack: The stack the pipeline will run on.
              base_environment: Base environment shared by all steps. This should
                  be set if your orchestrator for example runs one container that
                  is responsible for starting all the steps.
              step_environments: Environment variables to set when executing
                  specific steps.
              placeholder_run: An optional placeholder run for the snapshot.
  
          Returns:
              Optional submission result.
          """
          settings = cast(
              LightningOrchestratorSettings, self.get_settings(snapshot)
          )
          if snapshot.schedule:
              logger.warning(
                  "Lightning Orchestrator currently does not support the "
                  "use of schedules. The `schedule` will be ignored "
                  "and the pipeline will be run immediately."
              )
  
          snapshot_id = snapshot.id
  
          pipeline_name = snapshot.pipeline_configuration.name
          orchestrator_run_name = get_orchestrator_run_name(pipeline_name)
  
          code_archive = code_utils.CodeArchive(
              root=source_utils.get_source_root()
          )
          logger.info("Archiving pipeline code...")
          with tempfile.NamedTemporaryFile(
              mode="w+b", delete=False, suffix=".tar.gz"
          ) as code_file:
              code_archive.write_archive(code_file)
              code_path = code_file.name
          filename = f"{orchestrator_run_name}.tar.gz"
          # Construct the env variables for the pipeline
          env_vars = base_environment.copy()
          orchestrator_run_id = str(uuid4())
          env_vars[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID] = orchestrator_run_id
          # Set up some variables for configuration
          env_vars[ENV_ZENML_CUSTOM_SOURCE_ROOT] = (
              LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH
          )
  
          # Create a line-by-line export of environment variables
          env_exports = "\n".join(
              [f"export {key}='{value}'" for key, value in env_vars.items()]
          )
  
          # Write the environment variables to a temporary file
          with tempfile.NamedTemporaryFile(
              mode="w", delete=False, suffix=".studiorc"
          ) as temp_file:
              temp_file.write(env_exports)
              env_file_path = temp_file.name
  
          # Gather the requirements
          pipeline_docker_settings = (
              snapshot.pipeline_configuration.docker_settings
          )
          pipeline_requirements = gather_requirements(pipeline_docker_settings)
          pipeline_requirements_to_string = " ".join(
              f'"{req}"' for req in pipeline_requirements
          )
  
          def _construct_lightning_steps(
              snapshot: "PipelineSnapshotResponse",
          ) -> Dict[str, Dict[str, Any]]:
              """Construct the steps for the pipeline.
  
              Args:
                  snapshot: The pipeline snapshot to prepare or run.
  
              Returns:
                  The steps for the pipeline.
              """
              steps = {}
              for step_name, step in snapshot.step_configurations.items():
                  # The arguments are passed to configure the entrypoint of the
                  # docker container when the step is called.
                  entrypoint_command = (
                      StepEntrypointConfiguration.get_entrypoint_command()
                  )
                  entrypoint_arguments = (
                      StepEntrypointConfiguration.get_entrypoint_arguments(
                          step_name=step_name,
                          snapshot_id=snapshot_id,
                      )
                  )
                  entrypoint = entrypoint_command + entrypoint_arguments
                  entrypoint_string = " ".join(entrypoint)
  
                  step_settings = cast(
                      LightningOrchestratorSettings, self.get_settings(step)
                  )
  
                  # Gather the requirements
                  step_docker_settings = step.config.docker_settings
                  step_requirements = gather_requirements(step_docker_settings)
                  step_requirements_to_string = " ".join(
                      f'"{req}"' for req in step_requirements
                  )
  
                  # Construct the command to run the step
                  run_command = f"{entrypoint_string}"
                  commands = [run_command]
                  steps[step_name] = {
                      "commands": commands,
                      "requirements": step_requirements_to_string,
                      "machine": step_settings.machine_type
                      if step_settings != settings
                      else None,
                  }
              return steps
  
          if not settings.synchronous:
              entrypoint_command = LightningOrchestratorEntrypointConfiguration.get_entrypoint_command()
              entrypoint_arguments = LightningOrchestratorEntrypointConfiguration.get_entrypoint_arguments(
                  run_name=orchestrator_run_name,
                  snapshot_id=snapshot.id,
              )
              entrypoint = entrypoint_command + entrypoint_arguments
              entrypoint_string = " ".join(entrypoint)
              logger.info("Setting up Lightning AI client")
              self._set_lightning_env_vars(snapshot)
  
              studio_name = sanitize_studio_name(
                  "zenml_async_orchestrator_studio"
              )
              logger.info(f"Creating main studio: {studio_name}")
              studio = Studio(name=studio_name)
              studio.start()
  
              logger.info(
                  "Uploading wheel package and installing dependencies on main studio"
              )
              studio.run(
                  f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
              )
              studio.upload_file(
                  code_path,
                  remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
              )
              time.sleep(10)
              studio.run(
                  f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
              )
              studio.upload_file(env_file_path)
              time.sleep(6)
              studio.run(
                  f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
              )
              studio.run(f"rm {env_file_path.split('/')[-1]}")
  
              studio.run("pip install uv")
              logger.info(
                  f"Installing requirements: {pipeline_requirements_to_string}"
              )
              studio.run(f"uv pip install {pipeline_requirements_to_string}")
              studio.run("pip install zenml")
  
              for custom_command in settings.custom_commands or []:
                  studio.run(
                      f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {custom_command}"
                  )
  
              logger.info("Running pipeline in async mode")
              studio.run(
                  f"nohup bash -c 'cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {entrypoint_string}' > log_{filename.rsplit('.', 2)[0]}.txt 2>&1 &"
              )
              logger.info(
                  f"The pipeline is running in async mode, you can keep checking the logs by running the following command: `lightning download -s vision-model/zenml-async-orchestrator-studio -p /teamspace/studios/this_studio/log_{filename.rsplit('.', 2)[0]}.txt && cat log_{filename.rsplit('.', 2)[0]}.txt`"
              )
          else:
              self._upload_and_run_pipeline(
                  snapshot,
                  orchestrator_run_id,
                  pipeline_requirements_to_string,
                  settings,
                  _construct_lightning_steps(snapshot),
                  code_path,
                  filename,
                  env_file_path,
              )
          os.unlink(env_file_path)
          return None
  
      def _upload_and_run_pipeline(
          self,
          snapshot: "PipelineSnapshotResponse",
          orchestrator_run_id: str,
          requirements: str,
          settings: LightningOrchestratorSettings,
          steps_commands: Dict[str, Dict[str, Any]],
          code_path: str,
          filename: str,
          env_file_path: str,
      ) -> None:
          """Upload and run the pipeline on Lightning Studio.
  
          Args:
              snapshot: The pipeline snapshot to prepare or run.
              orchestrator_run_id: The orchestrator run id.
              requirements: The requirements for the pipeline.
              settings: The orchestrator settings.
              steps_commands: The commands for the steps.
              code_path: The path to the wheel package.
              filename: The name of the code archive.
              env_file_path: The path to the environment file.
  
          Raises:
              Exception: If an error occurs while running the pipeline.
          """
          logger.info("Setting up Lightning AI client")
          self._set_lightning_env_vars(snapshot)
  
          if settings.main_studio_name:
              studio_name = settings.main_studio_name
              studio = Studio(name=studio_name)
              if (
                  studio.machine != settings.machine_type
                  and settings.machine_type
              ):
                  studio.switch_machine(Machine(settings.machine_type))
          else:
              studio_name = sanitize_studio_name(
                  f"zenml_{orchestrator_run_id}_pipeline"
              )
              logger.info(f"Creating main studio: {studio_name}")
              studio = Studio(name=studio_name)
              if settings.machine_type:
                  studio.start(Machine(settings.machine_type))
              else:
                  studio.start()
          try:
              studio.run(
                  f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
              )
              studio.upload_file(
                  code_path,
                  remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
              )
              time.sleep(6)
              studio.run(
                  f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
              )
              logger.info(
                  "Uploading wheel package and installing dependencies on main studio"
              )
              # for some reason uploading file directly to /teamspace/studios/this_studio/.lightning_studio/.studiorc
              # doesn't work, while other file names can be uploaded just fine.
              # below is a workaround to copy the file to the correct location.
              # we first upload it under a different name and then copy it to the
              # right location.
              studio.upload_file(env_file_path)
              time.sleep(6)
              studio.run(
                  f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
              )
              studio.run(f"rm {env_file_path.split('/')[-1]}")
  
              studio.run("pip install uv")
              studio.run(f"uv pip install {requirements}")
              studio.run(f"uv pip install zenml=={__version__}")
              for command in settings.custom_commands or []:
                  output = studio.run(
                      f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
                  )
                  logger.info(f"Custom command output: {output}")
  
              for step_name, details in steps_commands.items():
                  if details["machine"]:
                      logger.info(f"Executing step: {step_name} in new studio")
                      self._run_step_in_new_studio(
                          orchestrator_run_id,
                          step_name,
                          details,
                          code_path,
                          filename,
                          env_file_path,
                          settings.custom_commands,
                      )
                  else:
                      logger.info(f"Executing step: {step_name} in main studio")
                      self._run_step_in_main_studio(studio, details, filename)
          except Exception as e:
              logger.error(f"Error running pipeline: {e}")
              raise e
          finally:
              if (
                  studio.status != studio.status.NotCreated
                  and settings.main_studio_name is None
              ):
                  logger.info("Deleting main studio")
                  studio.delete()
  
      def _run_step_in_new_studio(
          self,
          orchestrator_run_id: str,
          step_name: str,
          details: Dict[str, Any],
          code_path: str,
          filename: str,
          env_file_path: str,
          custom_commands: Optional[List[str]] = None,
      ) -> None:
          """Run a step in a new studio.
  
          Args:
              orchestrator_run_id: The orchestrator run id.
              step_name: The name of the step.
              details: The details of the step.
              code_path: The path to the wheel package.
              filename: The name of the code archive.
              env_file_path: The path to the environment file.
              custom_commands: Custom commands to run.
          """
          studio_name = sanitize_studio_name(
              f"zenml_{orchestrator_run_id}_{step_name}"
          )
          logger.info(f"Creating new studio for step {step_name}: {studio_name}")
          studio = Studio(name=studio_name)
          studio.start(Machine(details["machine"]))
          studio.run(
              f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
          )
          studio.upload_file(
              code_path,
              remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
          )
          time.sleep(6)
          studio.run(
              f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
          )
          studio.upload_file(env_file_path)
          time.sleep(6)
          studio.run(
              f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
          )
          studio.run(f"rm {env_file_path.split('/')[-1]}")
  
          studio.run("pip install uv")
          studio.run(f"uv pip install {details['requirements']}")
          studio.run(f"uv pip install zenml=={__version__}")
          for command in custom_commands or []:
              output = studio.run(
                  f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
              )
              logger.info(f"Custom command output: {output}")
          for command in details["commands"]:
              output = studio.run(
                  f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
              )
              logger.info(f"Step {step_name} output: {output}")
          studio.delete()
  
      def _run_step_in_main_studio(
          self, studio: Studio, details: Dict[str, Any], filename: str
      ) -> None:
          """Run a step in the main studio.
  
          Args:
              studio: The studio to run the step in.
              details: The details of the step.
              filename: The name of the code archive.
          """
          for command in details["commands"]:
              output = studio.run(
                  f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
              )
              logger.info(f"Step output: {output}")
  