





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_model_parallelism.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_600, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_600, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
    ]				)
class     UpperCamelCase__					(     unittest.TestCase				):






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> Any:






									"""simple docstring"""
									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=snake_case__,   )
									assert hasattr(self,   """env"""  )






		def 						snake_case__				(							self,   snake_case__  )       -> int:






									"""simple docstring"""
									# configuration for running training on smdistributed Model Parallel
									lowercase_	:   Dict  				=     {
									    """enabled""": True,
									    """processes_per_host""": 8,
									}
									lowercase_	:   Dict  				=     {
									    """enabled""": True,
									    """parameters""": {
									        """microbatches""": 4,
									        """placement_strategy""": """spread""",
									        """pipeline""": """interleaved""",
									        """optimize""": """speed""",
									        """partitions""": 4,
									        """ddp""": True,
									    },
									}

									lowercase_	:   List[str]  				=     {"""smdistributed""": {"""modelparallel""": smp_options}, """mpi""": mpi_options}

									lowercase_	:   List[Any]  				=     """trainer""" if self.script == """run_glue.py""" else """smtrainer"""
									# 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=f"""{self.env.base_job_name}-{instance_count}-smp-{name_extension}""",   instance_count=snake_case__,   instance_type=self.instance_type,   debugger_hook_config=snake_case__,   hyperparameters={
									        **self.env.hyperparameters,
									        """model_name_or_path""": self.model_name_or_path,
									        """max_steps""": 5_00,
									    },   metric_definitions=self.env.metric_definitions,   distribution=snake_case__,   py_version="""py36""",   )






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									TrainingJobAnalytics(snake_case__  ).export_csv(f"""{self.env.test_path}/{job_name}_metrics.csv"""  )






		@parameterized.expand([(1,)]  )
		def 						snake_case__				(							self,   snake_case__  )       -> List[Any]:






									"""simple docstring"""
									# create estimator
									lowercase_	:   Any  				=     self.create_estimator(snake_case__  )

									# run training
									estimator.fit()

									# result dataframe
									lowercase_	:   Any  				=     TrainingJobAnalytics(estimator.latest_training_job.name  ).dataframe()

									# extract kpis
									lowercase_	:   Optional[Any]  				=     list(result_metrics_df[result_metrics_df.metric_name == """eval_accuracy"""]["""value"""]  )
									lowercase_	:   Optional[Any]  				=     list(result_metrics_df[result_metrics_df.metric_name == """eval_loss"""]["""value"""]  )
									# get train time from SageMaker job, this includes starting, preprocessing, stopping
									lowercase_	:   Dict  				=     (
									    Session().describe_training_job(estimator.latest_training_job.name  ).get("""TrainingTimeInSeconds""",   99_99_99  )
									)

									# 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},   snake_case__  )






import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp





def  __magic_name__  (				)				-> str:

							"""simple docstring"""



							lowercase_	:   Optional[int]  				=     ArgumentParser(
							    description=(
							        """PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"""
							    )				)

							# Optional arguments for the launch helper
							parser.add_argument("""--num_cores"""				,       type=lowercase				,       default=1				,       help="""Number of TPU cores to use (1 or 8)."""				)

							# positional
							parser.add_argument(
							    """training_script"""				,       type=lowercase				,       help=(
							        """The full path to the single TPU training """
							        """program/script to be launched in parallel, """
							        """followed by all the arguments for the """
							        """training script"""
							    )				,       )

							# rest from the training program
							parser.add_argument("""training_script_args"""				,       nargs=lowercase				)

							return parser.parse_args()





def  __magic_name__  (				)				-> List[Any]:

							"""simple docstring"""



							lowercase_	:   Union[str, Any]  				=     parse_args()

							# Import training_script as a module.
							lowercase_	:   List[Any]  				=     Path(args.training_script				)
							sys.path.append(str(script_fpath.parent.resolve()				)				)
							lowercase_	:   Tuple  				=     script_fpath.stem
							lowercase_	:   List[str]  				=     importlib.import_module(lowercase				)

							# Patch sys.argv
							lowercase_	:   Optional[int]  				=     [args.training_script] + args.training_script_args + ["""--tpu_num_cores""", str(args.num_cores				)]

							xmp.spawn(mod._mp_fn				,       args=()				,       nprocs=args.num_cores				)


if __name__ == "__main__":
							main()