






'''simple docstring'''



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


SCREAMING_SNAKE_CASE							=						logging.get_logger(__name__)


def 				lowercase_					(     )							->   Tuple:



 """simple docstring"""




 lowercase       : Optional[Any]						=os.getenv('''SM_HP_MP_PARAMETERS'''       ,   '''{}'''      )
 try:
  # Parse it and check the field "partitions" is included, it is required for model parallel.
  lowercase       : List[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.
 lowercase       : Tuple						=os.getenv('''SM_FRAMEWORK_PARAMS'''       ,   '''{}'''      )
 try:
  # Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
  lowercase       : Dict						=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 UpperCAmelCase_       (     __A					):




      """simple docstring"""




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






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


       '''simple docstring'''

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






      @cached_property
      def 					A__   (			self      :		List[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:
        lowercase       : Optional[Any]						=torch.device('''cpu'''		)
        lowercase       : List[str]						=0
       elif is_sagemaker_model_parallel_available():
        lowercase       : List[Any]						=smp.local_rank()
        lowercase       : List[str]						=torch.device('''cuda'''  ,   UpperCAmelCase		)
        lowercase       : List[Any]						=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		)
        lowercase       : Dict						=int(os.getenv('''SMDATAPARALLEL_LOCAL_RANK'''		)		)
        lowercase       : Union[str, Any]						=torch.device('''cuda'''  ,   self.local_rank		)
        lowercase       : Optional[Any]						=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
        lowercase       : Optional[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.
        lowercase       : Dict						=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		)
        lowercase       : str						=torch.device('''cuda'''  ,   self.local_rank		)
        lowercase       : Union[str, Any]						=1

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

       return device






      @property
      def 					A__   (			self      :		Tuple		)	->	str:


       '''simple docstring'''

       if is_sagemaker_model_parallel_available():
        return smp.dp_size()

       return super().world_size






      @property
      def 					A__   (			self      :		Optional[int]		)	->	List[str]:


       '''simple docstring'''

       return not is_sagemaker_model_parallel_available()






      @property
      def 					A__   (			self      :		Tuple		)	->	Any:


       '''simple docstring'''

       return False











'''simple docstring'''



import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate


def 				lowercase_					(     __A			:							str      )							->   Union[str, Any]:



 """simple docstring"""




 return x + 2





class UpperCAmelCase_       (     unittest.TestCase					):




      """simple docstring"""




      def 					A__   (			self      :		Dict		)	->	List[Any]:


       '''simple docstring'''

       lowercase       : Optional[int]						='''x = 3'''
       lowercase       : Any						={}
       lowercase       : Union[str, Any]						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       assert result == 3
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3}		)

       lowercase       : str						='''x = y'''
       lowercase       : Optional[int]						={'''y''': 5}
       lowercase       : List[str]						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       # evaluate returns the value of the last assignment.
       assert result == 5
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 5, '''y''': 5}		)






      def 					A__   (			self      :		str		)	->	Optional[int]:


       '''simple docstring'''

       lowercase       : Optional[int]						='''y = add_two(x)'''
       lowercase       : str						={'''x''': 3}
       lowercase       : List[str]						=evaluate(UpperCAmelCase  ,   {'''add_two''': add_two}  ,   state=UpperCAmelCase		)
       assert result == 5
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''y''': 5}		)

       # Won't work without the tool
       with CaptureStdout() as out:
        lowercase       : Optional[Any]						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       assert result is None
       assert "tried to execute add_two" in out.out






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


       '''simple docstring'''

       lowercase       : int						='''x = 3'''
       lowercase       : Dict						={}
       lowercase       : List[Any]						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       assert result == 3
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3}		)






      def 					A__   (			self      :		str		)	->	Tuple:


       '''simple docstring'''

       lowercase       : Optional[Any]						='''test_dict = {\'x\': x, \'y\': add_two(x)}'''
       lowercase       : str						={'''x''': 3}
       lowercase       : Tuple						=evaluate(UpperCAmelCase  ,   {'''add_two''': add_two}  ,   state=UpperCAmelCase		)
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''y''': 5}		)
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}		)






      def 					A__   (			self      :		Union[str, Any]		)	->	Union[str, Any]:


       '''simple docstring'''

       lowercase       : Optional[Any]						='''x = 3\ny = 5'''
       lowercase       : int						={}
       lowercase       : List[str]						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       # evaluate returns the value of the last assignment.
       assert result == 5
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''y''': 5}		)






      def 					A__   (			self      :		Any		)	->	Tuple:


       '''simple docstring'''

       lowercase       : List[str]						='''text = f\'This is x: {x}.\''''
       lowercase       : Union[str, Any]						={'''x''': 3}
       lowercase       : Tuple						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       # evaluate returns the value of the last assignment.
       assert result == "This is x: 3."
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''text''': '''This is x: 3.'''}		)






      def 					A__   (			self      :		Union[str, Any]		)	->	List[Any]:


       '''simple docstring'''

       lowercase       : Tuple						='''if x <= 3:\n    y = 2\nelse:\n    y = 5'''
       lowercase       : Union[str, Any]						={'''x''': 3}
       lowercase       : Any						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       # evaluate returns the value of the last assignment.
       assert result == 2
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''y''': 2}		)

       lowercase       : Optional[Any]						={'''x''': 8}
       lowercase       : str						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       # evaluate returns the value of the last assignment.
       assert result == 5
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 8, '''y''': 5}		)






      def 					A__   (			self      :		str		)	->	Optional[Any]:


       '''simple docstring'''

       lowercase       : List[str]						='''test_list = [x, add_two(x)]'''
       lowercase       : Any						={'''x''': 3}
       lowercase       : Tuple						=evaluate(UpperCAmelCase  ,   {'''add_two''': add_two}  ,   state=UpperCAmelCase		)
       self.assertListEqual(UpperCAmelCase  ,   [3, 5]		)
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''test_list''': [3, 5]}		)






      def 					A__   (			self      :		Any		)	->	Tuple:


       '''simple docstring'''

       lowercase       : str						='''y = x'''
       lowercase       : Dict						={'''x''': 3}
       lowercase       : Tuple						=evaluate(UpperCAmelCase  ,   {}  ,   state=UpperCAmelCase		)
       assert result == 3
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''y''': 3}		)






      def 					A__   (			self      :		str		)	->	List[str]:


       '''simple docstring'''

       lowercase       : Any						='''test_list = [x, add_two(x)]\ntest_list[1]'''
       lowercase       : Any						={'''x''': 3}
       lowercase       : Union[str, Any]						=evaluate(UpperCAmelCase  ,   {'''add_two''': add_two}  ,   state=UpperCAmelCase		)
       assert result == 5
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''test_list''': [3, 5]}		)

       lowercase       : int						='''test_dict = {\'x\': x, \'y\': add_two(x)}\ntest_dict[\'y\']'''
       lowercase       : Union[str, Any]						={'''x''': 3}
       lowercase       : Tuple						=evaluate(UpperCAmelCase  ,   {'''add_two''': add_two}  ,   state=UpperCAmelCase		)
       assert result == 5
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}		)






      def 					A__   (			self      :		List[str]		)	->	Tuple:


       '''simple docstring'''

       lowercase       : Optional[int]						='''x = 0\nfor i in range(3):\n    x = i'''
       lowercase       : List[str]						={}
       lowercase       : Union[str, Any]						=evaluate(UpperCAmelCase  ,   {'''range''': range}  ,   state=UpperCAmelCase		)
       assert result == 2
       self.assertDictEqual(UpperCAmelCase  ,   {'''x''': 2, '''i''': 2}		)



