

# 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 			A_      (    _UpperCAmelCase    ):
						SCREAMING_SNAKE_CASE_: Optional[Any]  		=       botoa.client("iam"    )

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







def 			A_      (    _UpperCAmelCase    ):
						SCREAMING_SNAKE_CASE_: Optional[Any]  		=       botoa.client("iam"    )
						return iam_client.get_role(RoleName=lowercase_    )["Role"]["Arn"]







def 			A_      (    ):
						SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       _ask_options(
						    "How do you want to authorize?"				,	["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "]				,	lowercase_				,	)
						SCREAMING_SNAKE_CASE_: str  		=       None
						if credentials_configuration == 0:
												SCREAMING_SNAKE_CASE_: Any  		=       _ask_field("Enter your AWS Profile name: [default] "				,	default="default"    )
												SCREAMING_SNAKE_CASE_: Any  		=       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_: Optional[Any]  		=       _ask_field("AWS Access Key ID: "    )
												SCREAMING_SNAKE_CASE_: Any  		=       aws_access_key_id

												SCREAMING_SNAKE_CASE_: Optional[Any]  		=       _ask_field("AWS Secret Access Key: "    )
												SCREAMING_SNAKE_CASE_: List[Any]  		=       aws_secret_access_key

						SCREAMING_SNAKE_CASE_: Dict  		=       _ask_field("Enter your AWS Region: [us-east-1]"				,	default="us-east-1"    )
						SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       aws_region

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

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

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

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

						SCREAMING_SNAKE_CASE_: List[str]  		=       _ask_options(
						    "What is the distributed mode?"				,	["No distributed training", "Data parallelism"]				,	_convert_sagemaker_distributed_mode				,	)
						SCREAMING_SNAKE_CASE_: Any  		=       {}
						SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       _ask_field(
						    "Do you wish to optimize your script with torch dynamo?[yes/NO]:"				,	_convert_yes_no_to_bool				,	default=lowercase_				,	error_message="Please enter yes or no."				,	)
						if use_dynamo:
												SCREAMING_SNAKE_CASE_: Optional[int]  		=       """dynamo_"""
												SCREAMING_SNAKE_CASE_: Optional[Any]  		=       _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_: Any  		=       _ask_field(
												    "Do you want to customize the defaults sent to torch.compile? [yes/NO]: "				,	_convert_yes_no_to_bool				,	default=lowercase_				,	error_message="Please enter yes or no."				,	)

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

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

						SCREAMING_SNAKE_CASE_: List[str]  		=       _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=lowercase_				,	compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER				,	distributed_type=lowercase_				,	use_cpu=lowercase_				,	dynamo_config=lowercase_				,	eca_instance_type=lowercase_				,	profile=lowercase_				,	region=lowercase_				,	iam_role_name=lowercase_				,	mixed_precision=lowercase_				,	num_machines=lowercase_				,	sagemaker_inputs_file=lowercase_				,	sagemaker_metrics_file=lowercase_				,	)









import re
import time
from typing import Optional

import IPython.display as disp

from ..trainer_callback import TrainerCallback
from ..trainer_utils import IntervalStrategy, has_length

def 						UpperCamelCase					(lowercase_: str					)     ->      Dict:
	A__							:    int        =		int(lowercase_					)
	A__   ,   A__   ,   A__							:    Tuple        =		t // 3600, (t // 60) % 60, t % 60
	return f"""{h}:{m:02d}:{s:02d}""" if h != 0 else f"""{m:02d}:{s:02d}"""

def 						UpperCamelCase					(lowercase_: str    ,      lowercase_: Optional[Any]    ,      lowercase_: Union[str, Any]    ,      lowercase_: Tuple    ,      lowercase_: Any=300					)     ->      Optional[int]:
	# docstyle-ignore
	return f"""
    <div>
      {prefix}
      <progress value='{value}' max='{total}' style='width:{width}px; height:20px; vertical-align: middle;'></progress>
      {label}
    </div>
    """

def 						UpperCamelCase					(lowercase_: Tuple					)     ->      Optional[int]:
	A__							:    Tuple        =		"""<table border=\"1\" class=\"dataframe\">\n"""
	html_code += """  <thead>\n <tr style="text-align: left;">\n"""
	for i in items[0]:
		html_code += f"""      <th>{i}</th>\n"""
	html_code += "    </tr>\n  </thead>\n  <tbody>\n"
	for line in items[1:]:
		html_code += "    <tr>\n"
		for elt in line:
			A__							:    str        =		f"""{elt:.6f}""" if isinstance(lowercase_    ,      lowercase_					) else str(lowercase_					)
			html_code += f"""      <td>{elt}</td>\n"""
		html_code += "    </tr>\n"
	html_code += "  </tbody>\n</table><p>"
	return html_code



class     _a						:


					'''simple docstring'''






					UpperCAmelCase__:   str  							=       5
					UpperCAmelCase__:   int  							=       0.2






					def __init__( self    ,       A__    ,       A__ = None    ,       A__ = True    ,       A__ = None    ,       A__ = 300    ,       ):
						A__							:    Optional[int]        =		total
						A__							:    Tuple        =		"""""" if prefix is None else prefix
						A__							:    str        =		leave
						A__							:    str        =		parent
						A__							:    int        =		width
						A__							:    Dict        =		None
						A__							:    List[str]        =		None
						A__							:    Optional[int]        =		None






					def   __A			( self    ,       A__    ,       A__ = False    ,       A__ = None			):
						A__							:    Tuple        =		value
						if comment is not None:
							A__							:    Any        =		comment
						if self.last_value is None:
							A__							:    int        =		time.time()
							A__							:    Dict        =		value
							A__							:    int        =		None
							A__							:    int        =		self.warmup
							A__							:    str        =		1
							self.update_bar(A__			)
						elif value <= self.last_value and not force_update:
							return
						elif force_update or self.first_calls > 0 or value >= min(self.last_value + self.wait_for    ,       self.total			):
							if self.first_calls > 0:
								self.first_calls -= 1
							A__							:    Any        =		time.time()
							A__							:    str        =		current_time - self.start_time
							# We could have value = self.start_value if the update is called twixe with the same start value.
							if value > self.start_value:
								A__							:    Dict        =		self.elapsed_time / (value - self.start_value)
							else:
								A__							:    List[str]        =		None
							if value >= self.total:
								A__							:    Optional[Any]        =		self.total
								A__							:    List[Any]        =		None
								if not self.leave:
									self.close()
							elif self.average_time_per_item is not None:
								A__							:    List[Any]        =		self.average_time_per_item * (self.total - value)
							self.update_bar(A__			)
							A__							:    Any        =		value
							A__							:    List[str]        =		current_time
							if self.average_time_per_item is None:
								A__							:    str        =		1
							else:
								A__							:    Optional[Any]        =		max(int(self.update_every / self.average_time_per_item			)    ,       1			)






					def   __A			( self    ,       A__    ,       A__=None			):
						A__							:    Tuple        =		""" """ * (len(str(self.total			)			) - len(str(A__			)			)) + str(A__			)
						if self.elapsed_time is None:
							A__							:    Union[str, Any]        =		F"""[{spaced_value}/{self.total} : < :"""
						elif self.predicted_remaining is None:
							A__							:    Tuple        =		F"""[{spaced_value}/{self.total} {format_time(self.elapsed_time			)}"""
						else:
							A__							:    Optional[int]        =		(
							    F"""[{spaced_value}/{self.total} {format_time(self.elapsed_time			)} <"""
							    F""" {format_time(self.predicted_remaining			)}"""
							)
							self.label += F""", {1/self.average_time_per_item:.2f} it/s"""
						self.label += "]" if self.comment is None or len(self.comment			) == 0 else F""", {self.comment}]"""
						self.display()






					def   __A			( self			):
						A__							:    str        =		html_progress_bar(self.value    ,       self.total    ,       self.prefix    ,       self.label    ,       self.width			)
						if self.parent is not None:
							# If this is a child bar, the parent will take care of the display.
							self.parent.display()
							return
						if self.output is None:
							A__							:    str        =		disp.display(disp.HTML(self.html_code			)    ,       display_id=A__			)
						else:
							self.output.update(disp.HTML(self.html_code			)			)






					def   __A			( self			):
						if self.parent is None and self.output is not None:
							self.output.update(disp.HTML(""""""			)			)



class     _a						(__magic_name__				):


					'''simple docstring'''






					def __init__( self    ,       A__    ,       A__=None			):
						super().__init__(A__			)
						A__							:    Optional[Any]        =		None if column_names is None else [column_names]
						A__							:    Optional[Any]        =		None






					def   __A			( self			):
						A__							:    List[str]        =		html_progress_bar(self.value    ,       self.total    ,       self.prefix    ,       self.label    ,       self.width			)
						if self.inner_table is not None:
							self.html_code += text_to_html_table(self.inner_table			)
						if self.child_bar is not None:
							self.html_code += self.child_bar.html_code
						if self.output is None:
							A__							:    Optional[int]        =		disp.display(disp.HTML(self.html_code			)    ,       display_id=A__			)
						else:
							self.output.update(disp.HTML(self.html_code			)			)






					def   __A			( self    ,       A__			):
						if self.inner_table is None:
							A__							:    List[str]        =		[list(values.keys()			), list(values.values()			)]
						else:
							A__							:    Optional[Any]        =		self.inner_table[0]
							if len(self.inner_table			) == 1:
								# We give a chance to update the column names at the first iteration
								for key in values.keys():
									if key not in columns:
										columns.append(A__			)
								A__							:    Any        =		columns
							self.inner_table.append([values[c] for c in columns]			)






					def   __A			( self    ,       A__    ,       A__=None    ,       A__=300			):
						A__							:    Optional[Any]        =		NotebookProgressBar(A__    ,       prefix=A__    ,       parent=self    ,       width=A__			)
						return self.child_bar






					def   __A			( self			):
						A__							:    List[str]        =		None
						self.display()



class     _a						(__magic_name__				):


					'''simple docstring'''






					def __init__( self			):
						A__							:    int        =		None
						A__							:    List[str]        =		None
						A__							:    Union[str, Any]        =		False






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						A__							:    List[str]        =		"""Epoch""" if args.evaluation_strategy == IntervalStrategy.EPOCH else """Step"""
						A__							:    Dict        =		0
						A__							:    Tuple        =		0
						A__							:    Optional[int]        =		[self.first_column] + ["""Training Loss"""]
						if args.evaluation_strategy != IntervalStrategy.NO:
							column_names.append("""Validation Loss"""			)
						A__							:    Union[str, Any]        =		NotebookTrainingTracker(state.max_steps    ,       A__			)






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						A__							:    Any        =		int(state.epoch			) if int(state.epoch			) == state.epoch else F"""{state.epoch:.2f}"""
						self.training_tracker.update(
						    state.global_step + 1    ,       comment=F"""Epoch {epoch}/{state.num_train_epochs}"""    ,       force_update=self._force_next_update    ,       )
						A__							:    str        =		False






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__=None    ,       **A__			):
						if not has_length(A__			):
							return
						if self.prediction_bar is None:
							if self.training_tracker is not None:
								A__							:    Union[str, Any]        =		self.training_tracker.add_child(len(A__			)			)
							else:
								A__							:    Tuple        =		NotebookProgressBar(len(A__			)			)
							self.prediction_bar.update(1			)
						else:
							self.prediction_bar.update(self.prediction_bar.value + 1			)






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						if self.prediction_bar is not None:
							self.prediction_bar.close()
						A__							:    List[str]        =		None






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__=None    ,       **A__			):
						# Only for when there is no evaluation
						if args.evaluation_strategy == IntervalStrategy.NO and "loss" in logs:
							A__							:    Dict        =		{"""Training Loss""": logs["""loss"""]}
							# First column is necessarily Step sine we're not in epoch eval strategy
							A__							:    List[Any]        =		state.global_step
							self.training_tracker.write_line(A__			)






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__=None    ,       **A__			):
						if self.training_tracker is not None:
							A__							:    Tuple        =		{"""Training Loss""": """No log""", """Validation Loss""": """No log"""}
							for log in reversed(state.log_history			):
								if "loss" in log:
									A__							:    Dict        =		log["""loss"""]
									break

							if self.first_column == "Epoch":
								A__							:    List[Any]        =		int(state.epoch			)
							else:
								A__							:    Optional[Any]        =		state.global_step
							A__							:    Optional[Any]        =		"""eval"""
							for k in metrics:
								if k.endswith("""_loss"""			):
									A__							:    Optional[int]        =		re.sub(r"""\_loss$"""    ,       """"""    ,       A__			)
							A__							:    int        =		metrics.pop("""total_flos"""    ,       A__			)
							A__							:    int        =		metrics.pop("""epoch"""    ,       A__			)
							A__							:    Optional[int]        =		metrics.pop(F"""{metric_key_prefix}_runtime"""    ,       A__			)
							A__							:    Any        =		metrics.pop(F"""{metric_key_prefix}_samples_per_second"""    ,       A__			)
							A__							:    List[Any]        =		metrics.pop(F"""{metric_key_prefix}_steps_per_second"""    ,       A__			)
							A__							:    Optional[Any]        =		metrics.pop(F"""{metric_key_prefix}_jit_compilation_time"""    ,       A__			)
							for k, v in metrics.items():
								if k == F"""{metric_key_prefix}_loss""":
									A__							:    Any        =		v
								else:
									A__							:    Optional[Any]        =		k.split("""_"""			)
									A__							:    Any        =		""" """.join([part.capitalize() for part in splits[1:]]			)
									A__							:    List[str]        =		v
							self.training_tracker.write_line(A__			)
							self.training_tracker.remove_child()
							A__							:    Dict        =		None
							# Evaluation takes a long time so we should force the next update.
							A__							:    Union[str, Any]        =		True






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						self.training_tracker.update(
						    state.global_step    ,       comment=F"""Epoch {int(state.epoch			)}/{state.num_train_epochs}"""    ,       force_update=A__			)
						A__							:    Optional[int]        =		None



