
from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)




def  __magic_name__ (   __lowerCAmelCase	:						Tuple		,  __lowerCAmelCase	:						Optional[Any]							)			->   List[Any]:
				__lowerCamelCase    					=	[]
				for part_id in partition_order:
								__lowerCamelCase    					=	df.where(f'''SPARK_PARTITION_ID() = {part_id}'''							).collect()
								for row_idx, row in enumerate(__lowerCAmelCase							):
												expected_row_ids_and_row_dicts.append((f'''{part_id}_{row_idx}''', row.asDict())							)
				return expected_row_ids_and_row_dicts




@require_not_windows
@require_dill_gt_0_3_2
def  __magic_name__ (   )			->   List[Any]:
				__lowerCamelCase    					=	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
				__lowerCamelCase    					=	spark.range(100							).repartition(1							)
				__lowerCamelCase    					=	Spark(__lowerCAmelCase							)
				# The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
				# that each partition can hold 2 rows.
				spark_builder._repartition_df_if_needed(max_shard_size=16							)
				# Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
				assert spark_builder.df.rdd.getNumPartitions() == 50




@require_not_windows
@require_dill_gt_0_3_2
def  __magic_name__ (   )			->   Dict:
				__lowerCamelCase    					=	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
				__lowerCamelCase    					=	spark.range(10							).repartition(2							)
				__lowerCamelCase    					=	[1, 0]
				__lowerCamelCase    					=	_generate_iterable_examples(__lowerCAmelCase		,  __lowerCAmelCase							)  # Reverse the partitions.
				__lowerCamelCase    					=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase		,  __lowerCAmelCase							)

				for i, (row_id, row_dict) in enumerate(generate_fn()							):
								__lowerCamelCase					,    __lowerCamelCase    					=	expected_row_ids_and_row_dicts[i]
								assert row_id == expected_row_id
								assert row_dict == expected_row_dict




@require_not_windows
@require_dill_gt_0_3_2
def  __magic_name__ (   )			->   int:
				__lowerCamelCase    					=	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
				__lowerCamelCase    					=	spark.range(10							).repartition(1							)
				__lowerCamelCase    					=	SparkExamplesIterable(__lowerCAmelCase							)
				assert it.n_shards == 1
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase							):
								assert row_id == f'''0_{i}'''
								assert row_dict == {"id": i}




@require_not_windows
@require_dill_gt_0_3_2
def  __magic_name__ (   )			->   Tuple:
				__lowerCamelCase    					=	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
				__lowerCamelCase    					=	spark.range(30							).repartition(3							)
				# Mock the generator so that shuffle reverses the partition indices.
				with patch('''numpy.random.Generator'''							) as generator_mock:
								__lowerCamelCase    					=	lambda __lowerCAmelCase							: x.reverse()
								__lowerCamelCase    					=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase		,  [2, 1, 0]							)

								__lowerCamelCase    					=	SparkExamplesIterable(__lowerCAmelCase							).shuffle_data_sources(__lowerCAmelCase							)
								assert shuffled_it.n_shards == 3
								for i, (row_id, row_dict) in enumerate(__lowerCAmelCase							):
												__lowerCamelCase					,    __lowerCamelCase    					=	expected_row_ids_and_row_dicts[i]
												assert row_id == expected_row_id
												assert row_dict == expected_row_dict




@require_not_windows
@require_dill_gt_0_3_2
def  __magic_name__ (   )			->   Optional[Any]:
				__lowerCamelCase    					=	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
				__lowerCamelCase    					=	spark.range(20							).repartition(4							)

				# Partitions 0 and 2
				__lowerCamelCase    					=	SparkExamplesIterable(__lowerCAmelCase							).shard_data_sources(worker_id=0		,  num_workers=2							)
				assert shard_it_a.n_shards == 2
				__lowerCamelCase    					=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase		,  [0, 2]							)
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase							):
								__lowerCamelCase					,    __lowerCamelCase    					=	expected_row_ids_and_row_dicts_a[i]
								assert row_id == expected_row_id
								assert row_dict == expected_row_dict

				# Partitions 1 and 3
				__lowerCamelCase    					=	SparkExamplesIterable(__lowerCAmelCase							).shard_data_sources(worker_id=1		,  num_workers=2							)
				assert shard_it_a.n_shards == 2
				__lowerCamelCase    					=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase		,  [1, 3]							)
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase							):
								__lowerCamelCase					,    __lowerCamelCase    					=	expected_row_ids_and_row_dicts_a[i]
								assert row_id == expected_row_id
								assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def  __magic_name__ (   )			->   Dict:
				__lowerCamelCase    					=	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
				__lowerCamelCase    					=	spark.range(100							).repartition(1							)
				__lowerCamelCase    					=	Spark(__lowerCAmelCase							)
				# Choose a small max_shard_size for maximum partitioning.
				spark_builder._repartition_df_if_needed(max_shard_size=1							)
				# The new number of partitions should not be greater than the number of rows.
				assert spark_builder.df.rdd.getNumPartitions() == 100





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    lowerCAmelCase__     (	unittest.TestCase			):






		def  __A       (		self		:							Optional[int]  )     ->							List[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=SCREAMING_SNAKE_CASE__							,			)
						assert hasattr(self							,			'''env'''  )






		def  __A       (		self		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							Optional[Any]:
						# configuration for running training on smdistributed Model Parallel
						__lowerCamelCase    					=	{
						    '''enabled''': True,
						    '''processes_per_host''': 8,
						}
						__lowerCamelCase    					=	{
						    '''enabled''': True,
						    '''parameters''': {
						        '''microbatches''': 4,
						        '''placement_strategy''': '''spread''',
						        '''pipeline''': '''interleaved''',
						        '''optimize''': '''speed''',
						        '''partitions''': 4,
						        '''ddp''': True,
						    },
						}

						__lowerCamelCase    					=	{'''smdistributed''': {'''modelparallel''': smp_options}, '''mpi''': mpi_options}

						__lowerCamelCase    					=	'''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=SCREAMING_SNAKE_CASE__							,			instance_type=self.instance_type							,			debugger_hook_config=SCREAMING_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=SCREAMING_SNAKE_CASE__							,			py_version='''py36'''							,			)






		def  __A       (		self		:							List[Any]							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							List[str]:
						TrainingJobAnalytics(SCREAMING_SNAKE_CASE__  ).export_csv(f'''{self.env.test_path}/{job_name}_metrics.csv'''  )







		@parameterized.expand([(1,)]  )
		def  __A       (		self		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]  )     ->							Optional[int]:
						# create estimator
						__lowerCamelCase    					=	self.create_estimator(SCREAMING_SNAKE_CASE__  )

						# run training
						estimator.fit()

						# result dataframe
						__lowerCamelCase    					=	TrainingJobAnalytics(estimator.latest_training_job.name  ).dataframe()

						# extract kpis
						__lowerCamelCase    					=	list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value''']  )
						__lowerCamelCase    					=	list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value''']  )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__lowerCamelCase    					=	(
						    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}							,			SCREAMING_SNAKE_CASE__  )



