


import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
						from sagemaker import Session, TrainingJobAnalytics
						from sagemaker.huggingface import HuggingFace







@pytest.mark.skipif(
    literal_eval(os.getenv("""TEST_SAGEMAKER""",						"""False"""							)							) is not True,						reason="""Skipping test because should only be run when releasing minor transformers version""",						)
@pytest.mark.usefixtures("""sm_env"""							)
@parameterized_class(
    [
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.p3.16xlarge""",
            """results""": {"""train_runtime""": 650, """eval_accuracy""": 0.7, """eval_loss""": 0.6},
        },
        {
            """framework""": """pytorch""",
            """script""": """run_ddp.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.p3.16xlarge""",
            """results""": {"""train_runtime""": 600, """eval_accuracy""": 0.7, """eval_loss""": 0.6},
        },
        {
            """framework""": """tensorflow""",
            """script""": """run_tf_dist.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.p3.16xlarge""",
            """results""": {"""train_runtime""": 600, """eval_accuracy""": 0.6, """eval_loss""": 0.7},
        },
    ]							)
class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										if self.framework == "pytorch":
														subprocess.run(
														    f"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()   ,						encoding='''utf-8'''   ,						check=__lowercase   ,						)
										assert hasattr(self   ,						'''env'''							)




						def 						lowercase_ (				self   ,						__lowercase							)					->    Any:
										lowerCAmelCase_ :      Tuple    				=				f"""{self.env.base_job_name}-{instance_count}-{"ddp" if "ddp" in self.script else "smd"}"""
										# distributed data settings
										lowerCAmelCase_ :      List[Any]    				=				{'''smdistributed''': {'''dataparallel''': {'''enabled''': True}}} if self.script != '''run_ddp.py''' else None

										# creates estimator
										return HuggingFace(
										    entry_point=self.script   ,						source_dir=self.env.test_path   ,						role=self.env.role   ,						image_uri=self.env.image_uri   ,						base_job_name=__lowercase   ,						instance_count=__lowercase   ,						instance_type=self.instance_type   ,						debugger_hook_config=__lowercase   ,						hyperparameters={**self.env.distributed_hyperparameters, '''model_name_or_path''': self.model_name_or_path}   ,						metric_definitions=self.env.metric_definitions   ,						distribution=__lowercase   ,						py_version='''py36'''   ,						)




						def 						lowercase_ (				self   ,						__lowercase							)					->    Tuple:
										TrainingJobAnalytics(__lowercase							).export_csv(f"""{self.env.test_path}/{job_name}_metrics.csv"""							)




						@parameterized.expand([(2,)]							)
						def 						lowercase_ (				self   ,						__lowercase							)					->    Dict:
										# create estimator
										lowerCAmelCase_ :      int    				=				self.create_estimator(__lowercase							)

										# run training
										estimator.fit()

										# result dataframe
										lowerCAmelCase_ :      Optional[int]    				=				TrainingJobAnalytics(estimator.latest_training_job.name							).dataframe()

										# extract kpis
										lowerCAmelCase_ :      List[str]    				=				list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value''']							)
										lowerCAmelCase_ :      str    				=				list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value''']							)
										# get train time from SageMaker job, this includes starting, preprocessing, stopping
										lowerCAmelCase_ :      str    				=				(
										    Session().describe_training_job(estimator.latest_training_job.name							).get('''TrainingTimeInSeconds'''   ,						9_9_9_9_9_9							)
										)

										# assert kpis
										assert train_runtime <= self.results["train_runtime"]
										assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy							)
										assert all(t <= self.results['''eval_loss'''] for t in eval_loss							)

										# dump tests result into json file to share in PR
										with open(f"""{estimator.latest_training_job.name}.json"""   ,						'''w'''							) as outfile:
														json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}   ,						__lowercase							)



# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("""ignore""", category=UserWarning, module="""torch.optim.lr_scheduler""")







class 				snake_case__:







						'''simple docstring'''




						def __init__(				self   ,						__lowercase   ,						__lowercase   ,						__lowercase = True   ,						__lowercase = False							)					->    Tuple:
										lowerCAmelCase_ :      Optional[int]    				=				scheduler
										lowerCAmelCase_ :      Dict    				=				optimizers if isinstance(__lowercase   ,						(list, tuple)							) else [optimizers]
										lowerCAmelCase_ :      str    				=				split_batches
										lowerCAmelCase_ :      Any    				=				step_with_optimizer
										lowerCAmelCase_ :      Optional[Any]    				=				GradientState()




						def 						lowercase_ (				self   ,						*__lowercase   ,						**__lowercase							)					->    Any:
										if not self.step_with_optimizer:
														# No link between scheduler and optimizer -> just step
														self.scheduler.step(*__lowercase   ,						**__lowercase							)
														return

										# Otherwise, first make sure the optimizer was stepped.
										if not self.gradient_state.sync_gradients:
														if self.gradient_state.adjust_scheduler:
																		self.scheduler._step_count += 1
														return

										for opt in self.optimizers:
														if opt.step_was_skipped:
																		return
										if self.split_batches:
														# Split batches -> the training dataloader batch size is not changed so one step per training step
														self.scheduler.step(*__lowercase   ,						**__lowercase							)
										else:
														# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
														# num_processes steps per training step
														lowerCAmelCase_ :      Optional[Any]    				=				AcceleratorState().num_processes
														for _ in range(__lowercase							):
																		# Special case when using OneCycle and `drop_last` was not used
																		if hasattr(self.scheduler   ,						'''total_steps'''							):
																						if self.scheduler._step_count <= self.scheduler.total_steps:
																										self.scheduler.step(*__lowercase   ,						**__lowercase							)
																		else:
																						self.scheduler.step(*__lowercase   ,						**__lowercase							)




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return self.scheduler.get_last_lr()




						def 						lowercase_ (				self							)					->    List[str]:
										return self.scheduler.state_dict()




						def 						lowercase_ (				self   ,						__lowercase							)					->    int:
										self.scheduler.load_state_dict(__lowercase							)




						def 						lowercase_ (				self							)					->    Tuple:
										return self.scheduler.get_lr()




						def 						lowercase_ (				self   ,						*__lowercase   ,						**__lowercase							)					->    int:
										return self.scheduler.print_lr(*__lowercase   ,						**__lowercase							)