content: # Apache Software License 2.0
  #
  # Copyright (c) ZenML GmbH 2024. 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 function to run any ZenML step using Accelerate."""
  
  import functools
  from typing import Any, Callable, Dict, Optional, TypeVar, Union, cast
  
  import cloudpickle as pickle
  from accelerate.commands.launch import (
      launch_command,
      launch_command_parser,
  )
  
  from zenml import get_pipeline_context
  from zenml.logger import get_logger
  from zenml.steps import BaseStep
  from zenml.utils.function_utils import _cli_arg_name, create_cli_wrapped_script
  
  logger = get_logger(__name__)
  F = TypeVar("F", bound=Callable[..., Any])
  
  
  def run_with_accelerate(
      step_function_top_level: Optional[BaseStep] = None,
      **accelerate_launch_kwargs: Any,
  ) -> Union[Callable[[BaseStep], BaseStep], BaseStep]:
      """Run a function with accelerate.
  
      Accelerate package: https://huggingface.co/docs/accelerate/en/index
      Example:
          ```python
          from zenml import step, pipeline
          from zenml.integrations.hugginface.steps import run_with_accelerate
  
          @run_with_accelerate(num_processes=4, multi_gpu=True)
          @step
          def training_step(some_param: int, ...):
              # your training code is below
              ...
  
          @pipeline
          def training_pipeline(some_param: int, ...):
              training_step(some_param, ...)
          ```
  
      Args:
          step_function_top_level: The step function to run with accelerate [optional].
              Used in functional calls like `run_with_accelerate(some_func,foo=bar)()`.
          accelerate_launch_kwargs: A dictionary of arguments to pass along to the
              `accelerate launch` command, including hardware selection, resource
              allocation, and training paradigm options. Visit
              https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-launch
              for more details.
  
      Returns:
          The accelerate-enabled version of the step.
      """
  
      def _decorator(step_function: BaseStep) -> BaseStep:
          def _wrapper(
              entrypoint: F, accelerate_launch_kwargs: Dict[str, Any]
          ) -> F:
              @functools.wraps(entrypoint)
              def inner(*args: Any, **kwargs: Any) -> Any:
                  if args:
                      raise ValueError(
                          "Accelerated steps do not support positional arguments."
                      )
  
                  with create_cli_wrapped_script(
                      entrypoint, flavor="accelerate"
                  ) as (
                      script_path,
                      output_path,
                  ):
                      commands = [str(script_path.absolute())]
                      for k, v in kwargs.items():
                          k = _cli_arg_name(k)
                          if isinstance(v, bool):
                              if v:
                                  commands.append(f"--{k}")
                          elif type(v) in (list, tuple, set):
                              for each in v:
                                  commands += [f"--{k}", f"{each}"]
                          else:
                              commands += [f"--{k}", f"{v}"]
                      logger.debug(commands)
  
                      parser = launch_command_parser()
                      args = parser.parse_args(commands)
                      for k, v in accelerate_launch_kwargs.items():
                          if k in args:
                              setattr(args, k, v)
                          else:
                              logger.warning(
                                  f"You passed in `{k}` as an `accelerate launch` argument, but it was not accepted. "
                                  "Please check https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-launch "
                                  "to find out more about supported arguments and retry."
                              )
                      try:
                          launch_command(args)
                      except Exception as e:
                          logger.error(
                              "Accelerate training job failed... See error message for details."
                          )
                          raise RuntimeError(
                              "Accelerate training job failed."
                          ) from e
                      else:
                          logger.info(
                              "Accelerate training job finished successfully."
                          )
                          return pickle.load(open(output_path, "rb"))
  
              return cast(F, inner)
  
          try:
              get_pipeline_context()
          except RuntimeError:
              pass
          else:
              raise RuntimeError(
                  f"`{run_with_accelerate.__name__}` decorator cannot be used "
                  "in a functional way with steps, please apply decoration "
                  "directly to a step instead. This behavior will be also "
                  "allowed in future, but now it faces technical limitations.\n"
                  "Example (allowed):\n"
                  f"@{run_with_accelerate.__name__}(...)\n"
                  f"def {step_function.name}(...):\n"
                  "    ...\n"
                  "Example (not allowed):\n"
                  "def my_pipeline(...):\n"
                  f"    run_with_accelerate({step_function.name},...)(...)\n"
              )
  
          setattr(
              step_function, "unwrapped_entrypoint", step_function.entrypoint
          )
          setattr(
              step_function,
              "entrypoint",
              _wrapper(
                  step_function.entrypoint,
                  accelerate_launch_kwargs=accelerate_launch_kwargs,
              ),
          )
  
          return step_function
  
      if step_function_top_level:
          return _decorator(step_function_top_level)
      return _decorator
  