

'''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 	_lowerCAmelCase (						_UpperCamelCase						:					Tuple    )				->      List[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   =botoa.client('iam'    )

     _SCREAMING_SNAKE_CASE							   ={
         '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    )    )
          _SCREAMING_SNAKE_CASE							   ={
              '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 	_lowerCAmelCase (						_UpperCamelCase						:					Optional[int]    )				->      Optional[int]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   =botoa.client('iam'    )
     return iam_client.get_role(RoleName=_UpperCamelCase    )["Role"]["Arn"]







def 	_lowerCAmelCase (						)				->      Any:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   =_ask_options(
         'How do you want to authorize?'							,			['AWS Profile', 'Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) ']							,			_UpperCamelCase							,			)
     _SCREAMING_SNAKE_CASE							   =None
     if credentials_configuration == 0:
          _SCREAMING_SNAKE_CASE							   =_ask_field('Enter your AWS Profile name: [default] '							,			default='default'    )
          _SCREAMING_SNAKE_CASE							   =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`'    )
          _SCREAMING_SNAKE_CASE							   =_ask_field('AWS Access Key ID: '    )
          _SCREAMING_SNAKE_CASE							   =aws_access_key_id

          _SCREAMING_SNAKE_CASE							   =_ask_field('AWS Secret Access Key: '    )
          _SCREAMING_SNAKE_CASE							   =aws_secret_access_key

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

     _SCREAMING_SNAKE_CASE							   =_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:
          _SCREAMING_SNAKE_CASE							   =_ask_field('Enter your IAM role name: '    )
     else:
          _SCREAMING_SNAKE_CASE							   ='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    )

     _SCREAMING_SNAKE_CASE							   =_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.'							,			)
     _SCREAMING_SNAKE_CASE							   =None
     if is_custom_docker_image:
          _SCREAMING_SNAKE_CASE							   =_ask_field('Enter your Docker image: '							,			lambda _UpperCamelCase    : str(_UpperCamelCase    ).lower()    )

     _SCREAMING_SNAKE_CASE							   =_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.'							,			)
     _SCREAMING_SNAKE_CASE							   =None
     if is_sagemaker_inputs_enabled:
          _SCREAMING_SNAKE_CASE							   =_ask_field(
              'Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): '							,			lambda _UpperCamelCase    : str(_UpperCamelCase    ).lower()							,			)

     _SCREAMING_SNAKE_CASE							   =_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.'							,			)
     _SCREAMING_SNAKE_CASE							   =None
     if is_sagemaker_metrics_enabled:
          _SCREAMING_SNAKE_CASE							   =_ask_field(
              'Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): '							,			lambda _UpperCamelCase    : str(_UpperCamelCase    ).lower()							,			)

     _SCREAMING_SNAKE_CASE							   =_ask_options(
         'What is the distributed mode?'							,			['No distributed training', 'Data parallelism']							,			_convert_sagemaker_distributed_mode							,			)
     _SCREAMING_SNAKE_CASE							   ={}
     _SCREAMING_SNAKE_CASE							   =_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:
          _SCREAMING_SNAKE_CASE							   ='dynamo_'
          _SCREAMING_SNAKE_CASE							   =_ask_options(
              'Which dynamo backend would you like to use?'							,			[x.lower() for x in DYNAMO_BACKENDS]							,			_convert_dynamo_backend							,			default=2							,			)
          _SCREAMING_SNAKE_CASE							   =_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:
               _SCREAMING_SNAKE_CASE							   =_ask_options(
                   'Which mode do you want to use?'							,			_UpperCamelCase							,			lambda _UpperCamelCase    : TORCH_DYNAMO_MODES[int(_UpperCamelCase    )]							,			default='default'							,			)
               _SCREAMING_SNAKE_CASE							   =_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.'							,			)
               _SCREAMING_SNAKE_CASE							   =_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.'							,			)
     _SCREAMING_SNAKE_CASE							   ='Which EC2 instance type you want to use for your training?'
     if distributed_type != SageMakerDistributedType.NO:
          _SCREAMING_SNAKE_CASE							   =_ask_options(
              _UpperCamelCase							,			_UpperCamelCase							,			lambda _UpperCamelCase    : SAGEMAKER_PARALLEL_EC2_INSTANCES[int(_UpperCamelCase    )]    )
     else:
          eca_instance_query += "? [ml.p3.2xlarge]:"
          _SCREAMING_SNAKE_CASE							   =_ask_field(_UpperCamelCase							,			lambda _UpperCamelCase    : str(_UpperCamelCase    ).lower()							,			default='ml.p3.2xlarge'    )

     _SCREAMING_SNAKE_CASE							   =1
     if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
          _SCREAMING_SNAKE_CASE							   =_ask_field(
              'How many machines do you want use? [1]: '							,			_UpperCamelCase							,			default=1							,			)

     _SCREAMING_SNAKE_CASE							   =_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							,			)



'''simple docstring'''
from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching


def 	_lowerCAmelCase (						)				->      int:
     """simple docstring"""




     import os as original_os
     from os import path as original_path
     from os import rename as original_rename
     from os.path import dirname as original_dirname
     from os.path import join as original_join

     assert _test_patching.os is original_os
     assert _test_patching.path is original_path
     assert _test_patching.join is original_join

     assert _test_patching.renamed_os is original_os
     assert _test_patching.renamed_path is original_path
     assert _test_patching.renamed_join is original_join

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_mock__'
     with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
          # Every way to access os.path.join must be patched, and the rest must stay untouched

          # check os.path.join
          assert isinstance(_test_patching.os							,			_PatchedModuleObj    )
          assert isinstance(_test_patching.os.path							,			_PatchedModuleObj    )
          assert _test_patching.os.path.join is mock

          # check path.join
          assert isinstance(_test_patching.path							,			_PatchedModuleObj    )
          assert _test_patching.path.join is mock

          # check join
          assert _test_patching.join is mock

          # check that the other attributes are untouched
          assert _test_patching.os.rename is original_rename
          assert _test_patching.path.dirname is original_dirname
          assert _test_patching.os.path.dirname is original_dirname

          # Even renamed modules or objects must be patched

          # check renamed_os.path.join
          assert isinstance(_test_patching.renamed_os							,			_PatchedModuleObj    )
          assert isinstance(_test_patching.renamed_os.path							,			_PatchedModuleObj    )
          assert _test_patching.renamed_os.path.join is mock

          # check renamed_path.join
          assert isinstance(_test_patching.renamed_path							,			_PatchedModuleObj    )
          assert _test_patching.renamed_path.join is mock

          # check renamed_join
          assert _test_patching.renamed_join is mock

          # check that the other attributes are untouched
          assert _test_patching.renamed_os.rename is original_rename
          assert _test_patching.renamed_path.dirname is original_dirname
          assert _test_patching.renamed_os.path.dirname is original_dirname

     # check that everthing is back to normal when the patch is over

     assert _test_patching.os is original_os
     assert _test_patching.path is original_path
     assert _test_patching.join is original_join

     assert _test_patching.renamed_os is original_os
     assert _test_patching.renamed_path is original_path
     assert _test_patching.renamed_join is original_join


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     assert _test_patching.open is open

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_builtin_mock__'
     # _test_patching has "open" in its globals
     assert _test_patching.open is open
     with patch_submodule(_test_patching							,			'open'							,			_UpperCamelCase    ):
          assert _test_patching.open is mock

     # check that everthing is back to normal when the patch is over

     assert _test_patching.open is open


def 	_lowerCAmelCase (						)				->      str:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_missing_mock__'
     with patch_submodule(_test_patching							,			'pandas.read_csv'							,			_UpperCamelCase    ):
          pass


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_missing_builtin_mock__'
     # _test_patching doesn't have "len" in its globals
     assert getattr(_test_patching							,			'len'							,			_UpperCamelCase    ) is None
     with patch_submodule(_test_patching							,			'len'							,			_UpperCamelCase    ):
          assert _test_patching.len is mock
     assert _test_patching.len is len


def 	_lowerCAmelCase (						)				->      str:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_start_and_stop_mock__'
     _SCREAMING_SNAKE_CASE							   =patch_submodule(_test_patching							,			'open'							,			_UpperCamelCase    )
     assert _test_patching.open is open
     patch.start()
     assert _test_patching.open is mock
     patch.stop()
     assert _test_patching.open is open


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     from os import rename as original_rename
     from os.path import dirname as original_dirname
     from os.path import join as original_join

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_join__'
     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_dirname__'
     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_rename__'
     assert _test_patching.os.path.join is original_join
     assert _test_patching.os.path.dirname is original_dirname
     assert _test_patching.os.rename is original_rename

     with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
          with patch_submodule(_test_patching							,			'os.rename'							,			_UpperCamelCase    ):
               with patch_submodule(_test_patching							,			'os.path.dirname'							,			_UpperCamelCase    ):
                    assert _test_patching.os.path.join is mock_join
                    assert _test_patching.os.path.dirname is mock_dirname
                    assert _test_patching.os.rename is mock_rename

    # try another order
     with patch_submodule(_test_patching							,			'os.rename'							,			_UpperCamelCase    ):
          with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
               with patch_submodule(_test_patching							,			'os.path.dirname'							,			_UpperCamelCase    ):
                    assert _test_patching.os.path.join is mock_join
                    assert _test_patching.os.path.dirname is mock_dirname
                    assert _test_patching.os.rename is mock_rename

     assert _test_patching.os.path.join is original_join
     assert _test_patching.os.path.dirname is original_dirname
     assert _test_patching.os.rename is original_rename







def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_doesnt_exist_mock__'
     with patch_submodule(_test_patching							,			'__module_that_doesn_exist__.__attribute_that_doesn_exist__'							,			_UpperCamelCase    ):
          pass
     with patch_submodule(_test_patching							,			'os.__attribute_that_doesn_exist__'							,			_UpperCamelCase    ):
          pass
