

from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


__UpperCAmelCase								=			logging.get_logger(__name__)


if is_torch_available():
   import torch
   from torch import nn
   from torch.utils.data import Dataset

   from transformers import Trainer



   class SCREAMING_SNAKE_CASE     (  a_			):
     """simple docstring"""






     def __init__(      self			:   str	,							lowerCAmelCase			:   int = 1_01   )							->     Optional[Any]:






        """simple docstring"""







        __lowerCAmelCase  :      Tuple    =						length







     def __len__(      self			:   Any   )							->     List[str]:






        """simple docstring"""







        return self.length







     def __getitem__(      self			:   Any	,							lowerCAmelCase			:   List[str]   )							->     int:






        """simple docstring"""







        return i






   class SCREAMING_SNAKE_CASE     :
     """simple docstring"""






     def __call__(      self			:   Tuple	,							lowerCAmelCase			:   Optional[int]   )							->     Union[str, Any]:






        """simple docstring"""







        return {"input_ids": torch.tensor(lowerCAmelCase   ), "labels": torch.tensor(lowerCAmelCase   )}






   class SCREAMING_SNAKE_CASE     (  nn.Module			):
     """simple docstring"""






     def __init__(      self			:   Optional[int]   )							->     Optional[int]:






        """simple docstring"""







        super().__init__()
        # Add some (unused) params otherwise DDP will complain.
        __lowerCAmelCase  :      int    =						nn.Linear(1_20	,							80   )







     def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple	,							lowerCAmelCase			:   Any	,							lowerCAmelCase			:   int=None   )							->     Dict:






        """simple docstring"""







        if labels is not None:
           return torch.tensor(0.0	,							device=input_ids.device   ), input_ids
        else:
           return input_ids



class SCREAMING_SNAKE_CASE     (  a_			):
  """simple docstring"""






  @require_torch_neuroncore
  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]   )							->     Optional[int]:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						f'''--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        '''.split()
     __lowerCAmelCase  :      Any    =						self.get_auto_remove_tmp_dir()
     __lowerCAmelCase  :      List[str]    =						f'''--output_dir {output_dir}'''.split()
     __lowerCAmelCase  :      Tuple    =						["""torchrun"""] + distributed_args + args
     execute_subprocess_async(lowerCAmelCase	,							env=self.get_env()   )
     # successful return here == success - any errors would have caused an error in the sub-call






class SCREAMING_SNAKE_CASE     (  a_			):
  """simple docstring"""






  @require_torch_multi_gpu
  def 	SCREAMING_SNAKE_CASE   (      self			:   Any   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Tuple    =						f'''--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        '''.split()
     __lowerCAmelCase  :      str    =						self.get_auto_remove_tmp_dir()
     __lowerCAmelCase  :      Optional[Any]    =						f'''--output_dir {output_dir}'''.split()
     __lowerCAmelCase  :      List[Any]    =						["""torchrun"""] + distributed_args + args
     execute_subprocess_async(lowerCAmelCase	,							env=self.get_env()   )
     # successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
   # The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
   #
   # PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

   __UpperCAmelCase								=			HfArgumentParser((TrainingArguments,))
   __UpperCAmelCase								=			parser.parse_args_into_dataclasses()[0]

   logger.warning(
       F'Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, '
       F'distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}'
   )

   # Essentially, what we want to verify in the distributed case is that we get all samples back,
   # in the right order. (this is crucial for prediction for instance)
   for dataset_length in [101, 40, 7]:
      __UpperCAmelCase								=			DummyDataset(dataset_length)





      def      snake_case_				(__A					:						EvalPrediction					)      ->     Dict:
         __lowerCAmelCase  :      Optional[int]    =						list(range(len(__A					)					)					)
         __lowerCAmelCase  :      Optional[Any]    =						p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
         if not success and training_args.local_rank == 0:
            logger.warning(
                """Predictions and/or labels do not match expected results:\n  - predictions: """
                f'''{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}'''					)
         return {"success": success}

      __UpperCAmelCase								=			Trainer(
          model=DummyModel(),
          args=training_args,
          data_collator=DummyDataCollator(),
          eval_dataset=dataset,
          compute_metrics=compute_metrics,
      )
      __UpperCAmelCase								=			trainer.evaluate()
      logger.info(metrics)
      if metrics["eval_success"] is not True:
         logger.error(metrics)
         exit(1)

      __UpperCAmelCase								=			trainer.predict(dataset)
      logger.info(p.metrics)
      if p.metrics["test_success"] is not True:
         logger.error(p.metrics)
         exit(1)

      __UpperCAmelCase								=			2

      __UpperCAmelCase								=			trainer.evaluate()
      logger.info(metrics)
      if metrics["eval_success"] is not True:
         logger.error(metrics)
         exit(1)

      __UpperCAmelCase								=			trainer.predict(dataset)
      logger.info(p.metrics)
      if p.metrics["test_success"] is not True:
         logger.error(p.metrics)
         exit(1)
      __UpperCAmelCase								=			None





import importlib.util
import json
import os
import warnings
from dataclasses import dataclass, field

import torch

from ..training_args import TrainingArguments
from ..utils import cached_property, is_sagemaker_dp_enabled, logging


__UpperCAmelCase								=			logging.get_logger(__name__)





def      snake_case_				()      ->     str:
   # Get the sagemaker specific mp parameters from smp_options variable.
   __lowerCAmelCase  :      List[Any]    =						os.getenv("""SM_HP_MP_PARAMETERS""" ,       """{}"""					)
   try:
      # Parse it and check the field "partitions" is included, it is required for model parallel.
      __lowerCAmelCase  :      Any    =						json.loads(__A					)
      if "partitions" not in smp_options:
         return False
   except json.JSONDecodeError:
      return False

   # Get the sagemaker specific framework parameters from mpi_options variable.
   __lowerCAmelCase  :      Any    =						os.getenv("""SM_FRAMEWORK_PARAMS""" ,       """{}"""					)
   try:
      # Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
      __lowerCAmelCase  :      Union[str, Any]    =						json.loads(__A					)
      if not mpi_options.get("""sagemaker_mpi_enabled""" ,       __A					):
         return False
   except json.JSONDecodeError:
      return False
   # Lastly, check if the `smdistributed` module is present.
   return importlib.util.find_spec("""smdistributed"""					) is not None


if is_sagemaker_model_parallel_available():
   import smdistributed.modelparallel.torch as smp

   smp.init()



@dataclass
class SCREAMING_SNAKE_CASE     (  a_			):
  """simple docstring"""






  lowerCamelCase   :      str										=field(
      default=""       ,							metadata={"help": "Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer"}       ,							)







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]   )							->     Optional[int]:






     """simple docstring"""







     super().__post_init__()
     warnings.warn(
         """`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use """
         """`TrainingArguments` instead."""	,							lowerCAmelCase	,							)







  @cached_property
  def 	SCREAMING_SNAKE_CASE   (      self			:   Any   )							->     "torch.device":






     """simple docstring"""







     logger.info("""PyTorch: setting up devices"""   )
     if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1:
        logger.warning(
            """torch.distributed process group is initialized, but local_rank == -1. """
            """In order to use Torch DDP, launch your script with `python -m torch.distributed.launch"""   )
     if self.no_cuda:
        __lowerCAmelCase  :      Optional[Any]    =						torch.device("""cpu"""   )
        __lowerCAmelCase  :      Union[str, Any]    =						0
     elif is_sagemaker_model_parallel_available():
        __lowerCAmelCase  :      int    =						smp.local_rank()
        __lowerCAmelCase  :      str    =						torch.device("""cuda"""	,							lowerCAmelCase   )
        __lowerCAmelCase  :      Tuple    =						1
     elif is_sagemaker_dp_enabled():
        import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401

        torch.distributed.init_process_group(backend="""smddp"""	,							timeout=self.ddp_timeout_delta   )
        __lowerCAmelCase  :      List[str]    =						int(os.getenv("""SMDATAPARALLEL_LOCAL_RANK"""   )   )
        __lowerCAmelCase  :      Tuple    =						torch.device("""cuda"""	,							self.local_rank   )
        __lowerCAmelCase  :      Optional[int]    =						1
     elif self.local_rank == -1:
        # if n_gpu is > 1 we'll use nn.DataParallel.
        # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
        # Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will
        # trigger an error that a device index is missing. Index 0 takes into account the
        # GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0`
        # will use the first GPU in that env, i.e. GPU#1
        __lowerCAmelCase  :      List[Any]    =						torch.device("""cuda:0""" if torch.cuda.is_available() else """cpu"""   )
        # Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
        # the default value.
        __lowerCAmelCase  :      Optional[int]    =						torch.cuda.device_count()
     else:
        # Here, we'll use torch.distributed.
        # Initializes the distributed backend which will take care of synchronizing nodes/GPUs
        if not torch.distributed.is_initialized():
           torch.distributed.init_process_group(backend="""nccl"""	,							timeout=self.ddp_timeout_delta   )
        __lowerCAmelCase  :      List[str]    =						torch.device("""cuda"""	,							self.local_rank   )
        __lowerCAmelCase  :      int    =						1

     if device.type == "cuda":
        torch.cuda.set_device(lowerCAmelCase   )

     return device







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   int   )							->     Any:






     """simple docstring"""







     if is_sagemaker_model_parallel_available():
        return smp.dp_size()

     return super().world_size







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]   )							->     Tuple:






     """simple docstring"""







     return not is_sagemaker_model_parallel_available()







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   int   )							->     Tuple:






     """simple docstring"""







     return False


