




'''simple docstring'''


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 						SCREAMING_SNAKE_CASE       (      a_	:							Dict			):
		__a			=					botoa.client('iam'			)

		__a			=					{
		    '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=a_				,					AssumeRolePolicyDocument=json.dumps(a_				,					indent=2			)			)
				__a			=					{
				    '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=a_				,					PolicyName=f"{role_name}_policy_permission"				,					PolicyDocument=json.dumps(a_				,					indent=2			)				,					)
		except iam_client.exceptions.EntityAlreadyExistsException:
				print(f"role {role_name} already exists. Using existing one"			)


def 						SCREAMING_SNAKE_CASE       (      a_	:							Union[str, Any]			):
		__a			=					botoa.client('iam'			)
		return iam_client.get_role(RoleName=a_			)["Role"]["Arn"]


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

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

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

		__a			=					_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']				,					a_				,					)
		if role_management == 0:
				__a			=					_ask_field('Enter your IAM role name: '			)
		else:
				__a			=					'accelerate_sagemaker_execution_role'
				print(f"Accelerate will create an iam role \"{iam_role_name}\" using the provided credentials"			)
				_create_iam_role_for_sagemaker(a_			)

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

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

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

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

				if use_custom_options:
						__a			=					_ask_options(
						    'Which mode do you want to use?'				,					a_				,					lambda a_			: TORCH_DYNAMO_MODES[int(a_			)]				,					default='default'				,					)
						__a			=					_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=a_				,					error_message='Please enter yes or no.'				,					)
						__a			=					_ask_field(
						    'Do you want to enable dynamic shape tracing? [yes/NO]: '				,					_convert_yes_no_to_bool				,					default=a_				,					error_message='Please enter yes or no.'				,					)
		__a			=					'Which EC2 instance type you want to use for your training?'
		if distributed_type != SageMakerDistributedType.NO:
				__a			=					_ask_options(
				    a_				,					a_				,					lambda a_			: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(a_			)]			)
		else:
				eca_instance_query += "? [ml.p3.2xlarge]:"
				__a			=					_ask_field(a_				,					lambda a_			: str(a_			).lower()				,					default='ml.p3.2xlarge'			)

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

		__a			=					_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=a_				,					compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER				,					distributed_type=a_				,					use_cpu=a_				,					dynamo_config=a_				,					eca_instance_type=a_				,					profile=a_				,					region=a_				,					iam_role_name=a_				,					mixed_precision=a_				,					num_machines=a_				,					sagemaker_inputs_file=a_				,					sagemaker_metrics_file=a_				,					)









'''simple docstring'''


from ..utils import DummyObject, requires_backends





class    __lowercase						(  metaclass=__magic_name__ ):
   _a							=      ["""onnx"""]







   def __init__(	self      ,							*UpperCamelCase      ,							**UpperCamelCase      )				->    str:
     requires_backends(self      ,							['onnx']      )







   @classmethod
   def 		UpperCamelCase__			(	cls      ,							*UpperCamelCase      ,							**UpperCamelCase      )				->    Tuple:
     requires_backends(cls      ,							['onnx']      )







   @classmethod
   def 		UpperCamelCase__			(	cls      ,							*UpperCamelCase      ,							**UpperCamelCase      )				->    Tuple:
     requires_backends(cls      ,							['onnx']      )



