






'''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 	__magic_name__( _A					):
						'''simple docstring'''





						UpperCamelCase__         = botoa.client("""iam"""					)

						UpperCamelCase__         = {
						    """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					)					)
												UpperCamelCase__         = {
												    """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 	__magic_name__( _A					):
						'''simple docstring'''





						UpperCamelCase__         = botoa.client("""iam"""					)
						return iam_client.get_role(RoleName=_A					)["Role"]["Arn"]







def 	__magic_name__( ):
						'''simple docstring'''





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

												UpperCamelCase__         = _ask_field("""AWS Secret Access Key: """					)
												UpperCamelCase__         = aws_secret_access_key

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

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

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

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

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

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

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

						UpperCamelCase__         = _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'''


import logging
import os
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


lowerCamelCase_				: Any    	=							{
    '''debug''': logging.DEBUG,
    '''info''': logging.INFO,
    '''warning''': logging.WARNING,
    '''error''': logging.ERROR,
    '''critical''': logging.CRITICAL,
}

lowerCamelCase_				: Optional[int]    	=							logging.WARNING







def 	__magic_name__( ):
						'''simple docstring'''





						UpperCamelCase__         = os.getenv("""DATASETS_VERBOSITY"""						,      _A					)
						if env_level_str:
												if env_level_str in log_levels:
																		return log_levels[env_level_str]
												else:
																		logging.getLogger().warning(
																		    f"Unknown option DATASETS_VERBOSITY={env_level_str}, "
																		    f"has to be one of: { ', '.join(log_levels.keys()					) }"					)
						return _default_log_level







def 	__magic_name__( ):
						'''simple docstring'''





						return __name__.split("""."""					)[0]







def 	__magic_name__( ):
						'''simple docstring'''





						return logging.getLogger(_get_library_name()					)







def 	__magic_name__( ):
						'''simple docstring'''





						UpperCamelCase__         = _get_library_root_logger()
						library_root_logger.setLevel(_get_default_logging_level()					)







def 	__magic_name__( ):
						'''simple docstring'''





						UpperCamelCase__         = _get_library_root_logger()
						library_root_logger.setLevel(logging.NOTSET					)







def 	__magic_name__( _A = None					):
						'''simple docstring'''





						if name is None:
												UpperCamelCase__         = _get_library_name()
						return logging.getLogger(_A					)







def 	__magic_name__( ):
						'''simple docstring'''





						return _get_library_root_logger().getEffectiveLevel()







def 	__magic_name__( _A					):
						'''simple docstring'''





						_get_library_root_logger().setLevel(_A					)







def 	__magic_name__( ):
						'''simple docstring'''





						return set_verbosity(_A					)







def 	__magic_name__( ):
						'''simple docstring'''





						return set_verbosity(_A					)







def 	__magic_name__( ):
						'''simple docstring'''





						return set_verbosity(_A					)







def 	__magic_name__( ):
						'''simple docstring'''





						return set_verbosity(_A					)







def 	__magic_name__( ):
						'''simple docstring'''





						UpperCamelCase__         = False







def 	__magic_name__( ):
						'''simple docstring'''





						UpperCamelCase__         = True


# Configure the library root logger at the module level (singleton-like)
_configure_library_root_logger()
class    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





			def __init__(   self							:						int ,      *lowercase							:						Any ,      **lowercase							:						int			)							->    List[Any]:  # pylint: disable=unused-argument

									'''simple docstring'''







									UpperCamelCase__         = args[0] if args else None






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

									'''simple docstring'''







									return iter(self._iterator			)






			def __getattr__(   self							:						Optional[int] ,      lowercase							:						Tuple			)							->    int:

									'''simple docstring'''







									def empty_fn(*lowercase							:						List[Any] ,      **lowercase							:						Dict			):  # pylint: disable=unused-argument
															return

									return empty_fn






			def __enter__(   self							:						Dict			)							->    Tuple:

									'''simple docstring'''







									return self







			def __exit__(   self							:						Tuple ,      lowercase							:						Optional[Any] ,      lowercase							:						Any ,      lowercase							:						int			)							->    Any:

									'''simple docstring'''







									return
lowerCamelCase_				: str    	=							True
class    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





			def __call__(   self							:						List[Any] ,      *lowercase							:						str ,      lowercase							:						Optional[int]=False ,      **lowercase							:						Optional[int]			)							->    Optional[int]:

									'''simple docstring'''







									if _tqdm_active and not disable:
															return tqdm_lib.tqdm(*lowercase ,      **lowercase			)
									else:
															return EmptyTqdm(*lowercase ,      **lowercase			)






			def 				A      (   self							:						Optional[Any] ,      *lowercase							:						Tuple ,      **lowercase							:						Any			)							->    Any:

									'''simple docstring'''







									UpperCamelCase__         = None
									if _tqdm_active:
															return tqdm_lib.tqdm.set_lock(*lowercase ,      **lowercase			)







			def 				A      (   self							:						Optional[Any]			)							->    Optional[Any]:

									'''simple docstring'''







									if _tqdm_active:
															return tqdm_lib.tqdm.get_lock()
lowerCamelCase_				: int    	=							_tqdm_cls()







def 	__magic_name__( ):
						'''simple docstring'''





						global _tqdm_active
						return bool(_tqdm_active					)







def 	__magic_name__( ):
						'''simple docstring'''





						global _tqdm_active
						UpperCamelCase__         = True







def 	__magic_name__( ):
						'''simple docstring'''





						global _tqdm_active
						UpperCamelCase__         = False
