'''simple docstring'''

# Copyright 2021 The HuggingFace Team. 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.
import json
import os

from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES
from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
from ...utils.imports import is_botoa_available
from .config_args import SageMakerConfig
from .config_utils import (
    DYNAMO_BACKENDS,
    _ask_field,
    _ask_options,
    _convert_dynamo_backend,
    _convert_mixed_precision,
    _convert_sagemaker_distributed_mode,
    _convert_yes_no_to_bool,
)


if is_botoa_available():
  import botoa  # noqa: F401







def __UpperCamelCase( _A       :					str     ):
   '''simple docstring'''
   UpperCAmelCase__				:					List[Any] 					=	botoa.client('''iam'''     )

   UpperCAmelCase__				:					str 					=	{
       '''Version''': '''2012-10-17''',
       '''Statement''': [
           {'''Effect''': '''Allow''', '''Principal''': {'''Service''': '''sagemaker.amazonaws.com'''}, '''Action''': '''sts:AssumeRole'''}
       ],
   }
   try:
      # create the role, associated with the chosen trust policy
      iam_client.create_role(
          RoleName=UpperCamelCase_	,      AssumeRolePolicyDocument=json.dumps(UpperCamelCase_	,      indent=2     )     )
      UpperCAmelCase__				:					List[str] 					=	{
          '''Version''': '''2012-10-17''',
          '''Statement''': [
              {
                  '''Effect''': '''Allow''',
                  '''Action''': [
                      '''sagemaker:*''',
                      '''ecr:GetDownloadUrlForLayer''',
                      '''ecr:BatchGetImage''',
                      '''ecr:BatchCheckLayerAvailability''',
                      '''ecr:GetAuthorizationToken''',
                      '''cloudwatch:PutMetricData''',
                      '''cloudwatch:GetMetricData''',
                      '''cloudwatch:GetMetricStatistics''',
                      '''cloudwatch:ListMetrics''',
                      '''logs:CreateLogGroup''',
                      '''logs:CreateLogStream''',
                      '''logs:DescribeLogStreams''',
                      '''logs:PutLogEvents''',
                      '''logs:GetLogEvents''',
                      '''s3:CreateBucket''',
                      '''s3:ListBucket''',
                      '''s3:GetBucketLocation''',
                      '''s3:GetObject''',
                      '''s3:PutObject''',
                  ],
                  '''Resource''': '''*''',
              }
          ],
      }
      # attach policy to role
      iam_client.put_role_policy(
          RoleName=UpperCamelCase_	,      PolicyName=F'''{role_name}_policy_permission'''	,      PolicyDocument=json.dumps(UpperCamelCase_	,      indent=2     )	,      )
   except iam_client.exceptions.EntityAlreadyExistsException:
      print(F'''role {role_name} already exists. Using existing one'''     )







def __UpperCamelCase( _A       :					List[Any]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					List[str] 					=	botoa.client('''iam'''     )
   return iam_client.get_role(RoleName=UpperCamelCase_     )["Role"]["Arn"]







def __UpperCamelCase( ):
   '''simple docstring'''
   UpperCAmelCase__				:					str 					=	_ask_options(
       '''How do you want to authorize?'''	,      ['''AWS Profile''', '''Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) ''']	,      UpperCamelCase_	,      )
   UpperCAmelCase__				:					Optional[int] 					=	None
   if credentials_configuration == 0:
      UpperCAmelCase__				:					Any 					=	_ask_field('''Enter your AWS Profile name: [default] '''	,      default='''default'''     )
      UpperCAmelCase__				:					str 					=	aws_profile
   else:
      print(
          '''Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with,'''
          '''`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`'''     )
      UpperCAmelCase__				:					Optional[int] 					=	_ask_field('''AWS Access Key ID: '''     )
      UpperCAmelCase__				:					Any 					=	aws_access_key_id

      UpperCAmelCase__				:					Any 					=	_ask_field('''AWS Secret Access Key: '''     )
      UpperCAmelCase__				:					str 					=	aws_secret_access_key

   UpperCAmelCase__				:					Optional[Any] 					=	_ask_field('''Enter your AWS Region: [us-east-1]'''	,      default='''us-east-1'''     )
   UpperCAmelCase__				:					Any 					=	aws_region

   UpperCAmelCase__				:					Any 					=	_ask_options(
       '''Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?'''	,      ['''Provide IAM Role name''', '''Create new IAM role using credentials''']	,      UpperCamelCase_	,      )
   if role_management == 0:
      UpperCAmelCase__				:					int 					=	_ask_field('''Enter your IAM role name: '''     )
   else:
      UpperCAmelCase__				:					List[Any] 					=	'''accelerate_sagemaker_execution_role'''
      print(F'''Accelerate will create an iam role "{iam_role_name}" using the provided credentials'''     )
      _create_iam_role_for_sagemaker(UpperCamelCase_     )

   UpperCAmelCase__				:					Tuple 					=	_ask_field(
       '''Do you want to use custom Docker image? [yes/NO]: '''	,      _convert_yes_no_to_bool	,      default=UpperCamelCase_	,      error_message='''Please enter yes or no.'''	,      )
   UpperCAmelCase__				:					List[str] 					=	None
   if is_custom_docker_image:
      UpperCAmelCase__				:					List[str] 					=	_ask_field('''Enter your Docker image: '''	,      lambda _A     : str(UpperCamelCase_     ).lower()     )

   UpperCAmelCase__				:					str 					=	_ask_field(
       '''Do you want to provide SageMaker input channels with data locations? [yes/NO]: '''	,      _convert_yes_no_to_bool	,      default=UpperCamelCase_	,      error_message='''Please enter yes or no.'''	,      )
   UpperCAmelCase__				:					List[Any] 					=	None
   if is_sagemaker_inputs_enabled:
      UpperCAmelCase__				:					List[Any] 					=	_ask_field(
          '''Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): '''	,      lambda _A     : str(UpperCamelCase_     ).lower()	,      )

   UpperCAmelCase__				:					str 					=	_ask_field(
       '''Do you want to enable SageMaker metrics? [yes/NO]: '''	,      _convert_yes_no_to_bool	,      default=UpperCamelCase_	,      error_message='''Please enter yes or no.'''	,      )
   UpperCAmelCase__				:					List[str] 					=	None
   if is_sagemaker_metrics_enabled:
      UpperCAmelCase__				:					Optional[Any] 					=	_ask_field(
          '''Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): '''	,      lambda _A     : str(UpperCamelCase_     ).lower()	,      )

   UpperCAmelCase__				:					str 					=	_ask_options(
       '''What is the distributed mode?'''	,      ['''No distributed training''', '''Data parallelism''']	,      _convert_sagemaker_distributed_mode	,      )
   UpperCAmelCase__				:					Optional[Any] 					=	{}
   UpperCAmelCase__				:					Tuple 					=	_ask_field(
       '''Do you wish to optimize your script with torch dynamo?[yes/NO]:'''	,      _convert_yes_no_to_bool	,      default=UpperCamelCase_	,      error_message='''Please enter yes or no.'''	,      )
   if use_dynamo:
      UpperCAmelCase__				:					Optional[int] 					=	'''dynamo_'''
      UpperCAmelCase__				:					int 					=	_ask_options(
          '''Which dynamo backend would you like to use?'''	,      [x.lower() for x in DYNAMO_BACKENDS]	,      _convert_dynamo_backend	,      default=2	,      )
      UpperCAmelCase__				:					Tuple 					=	_ask_field(
          '''Do you want to customize the defaults sent to torch.compile? [yes/NO]: '''	,      _convert_yes_no_to_bool	,      default=UpperCamelCase_	,      error_message='''Please enter yes or no.'''	,      )

      if use_custom_options:
         UpperCAmelCase__				:					List[str] 					=	_ask_options(
             '''Which mode do you want to use?'''	,      UpperCamelCase_	,      lambda _A     : TORCH_DYNAMO_MODES[int(UpperCamelCase_     )]	,      default='''default'''	,      )
         UpperCAmelCase__				:					int 					=	_ask_field(
             '''Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: '''	,      _convert_yes_no_to_bool	,      default=UpperCamelCase_	,      error_message='''Please enter yes or no.'''	,      )
         UpperCAmelCase__				:					Optional[Any] 					=	_ask_field(
             '''Do you want to enable dynamic shape tracing? [yes/NO]: '''	,      _convert_yes_no_to_bool	,      default=UpperCamelCase_	,      error_message='''Please enter yes or no.'''	,      )
   UpperCAmelCase__				:					Tuple 					=	'''Which EC2 instance type you want to use for your training?'''
   if distributed_type != SageMakerDistributedType.NO:
      UpperCAmelCase__				:					Optional[int] 					=	_ask_options(
          UpperCamelCase_	,      UpperCamelCase_	,      lambda _A     : SAGEMAKER_PARALLEL_EC2_INSTANCES[int(UpperCamelCase_     )]     )
   else:
      eca_instance_query += "? [ml.p3.2xlarge]:"
      UpperCAmelCase__				:					Optional[Any] 					=	_ask_field(UpperCamelCase_	,      lambda _A     : str(UpperCamelCase_     ).lower()	,      default='''ml.p3.2xlarge'''     )

   UpperCAmelCase__				:					Optional[int] 					=	1
   if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
      UpperCAmelCase__				:					Optional[Any] 					=	_ask_field(
          '''How many machines do you want use? [1]: '''	,      UpperCamelCase_	,      default=1	,      )

   UpperCAmelCase__				:					int 					=	_ask_options(
       '''Do you wish to use FP16 or BF16 (mixed precision)?'''	,      ['''no''', '''fp16''', '''bf16''', '''fp8''']	,      _convert_mixed_precision	,      )

   if use_dynamo and mixed_precision == "no":
      print(
          '''Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts.'''     )

   return SageMakerConfig(
       image_uri=UpperCamelCase_	,      compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER	,      distributed_type=UpperCamelCase_	,      use_cpu=UpperCamelCase_	,      dynamo_config=UpperCamelCase_	,      eca_instance_type=UpperCamelCase_	,      profile=UpperCamelCase_	,      region=UpperCamelCase_	,      iam_role_name=UpperCamelCase_	,      mixed_precision=UpperCamelCase_	,      num_machines=UpperCamelCase_	,      sagemaker_inputs_file=UpperCamelCase_	,      sagemaker_metrics_file=UpperCamelCase_	,      )








from ..utils import DummyObject, requires_backends

class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )






class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )






class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )






class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




