






'''simple docstring'''




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 	__UpperCAmelCase       ( __magic_name__     ,__magic_name__						)->    List[Any]:



			"""simple docstring"""



			snake_case_				:					str   =       []
			for part_id in partition_order:
						snake_case_				:					Dict   =       df.where(F'''SPARK_PARTITION_ID() = {part_id}'''						).collect()
						for row_idx, row in enumerate(__magic_name__						):
									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 	__UpperCAmelCase       ( )->    Tuple:



			"""simple docstring"""



			snake_case_				:					Dict   =       pyspark.sql.SparkSession.builder.master("local[*]"						).appName("pyspark"						).getOrCreate()
			snake_case_				:					List[Any]   =       spark.range(100						).repartition(1						)
			snake_case_				:					Tuple   =       Spark(__magic_name__						)
			# 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 	__UpperCAmelCase       ( )->    Optional[int]:



			"""simple docstring"""



			snake_case_				:					Union[str, Any]   =       pyspark.sql.SparkSession.builder.master("local[*]"						).appName("pyspark"						).getOrCreate()
			snake_case_				:					Dict   =       spark.range(10						).repartition(2						)
			snake_case_				:					Union[str, Any]   =       [1, 0]
			snake_case_				:					Optional[Any]   =       _generate_iterable_examples(__magic_name__     ,__magic_name__						)  # Reverse the partitions.
			snake_case_				:					Dict   =       _get_expected_row_ids_and_row_dicts_for_partition_order(__magic_name__     ,__magic_name__						)

			for i, (row_id, row_dict) in enumerate(generate_fn()						):
						snake_case_,	snake_case_				:					List[Any]   =       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 	__UpperCAmelCase       ( )->    Optional[Any]:



			"""simple docstring"""



			snake_case_				:					Optional[int]   =       pyspark.sql.SparkSession.builder.master("local[*]"						).appName("pyspark"						).getOrCreate()
			snake_case_				:					Optional[Any]   =       spark.range(10						).repartition(1						)
			snake_case_				:					str   =       SparkExamplesIterable(__magic_name__						)
			assert it.n_shards == 1
			for i, (row_id, row_dict) in enumerate(__magic_name__						):
						assert row_id == F'''0_{i}'''
						assert row_dict == {"id": i}

@require_not_windows
@require_dill_gt_0_3_2
def 	__UpperCAmelCase       ( )->    Union[str, Any]:



			"""simple docstring"""



			snake_case_				:					Any   =       pyspark.sql.SparkSession.builder.master("local[*]"						).appName("pyspark"						).getOrCreate()
			snake_case_				:					int   =       spark.range(30						).repartition(3						)
			# Mock the generator so that shuffle reverses the partition indices.
			with patch("numpy.random.Generator"						) as generator_mock:
						snake_case_				:					List[str]   =       lambda __magic_name__						: x.reverse()
						snake_case_				:					Dict   =       _get_expected_row_ids_and_row_dicts_for_partition_order(__magic_name__     ,[2, 1, 0]						)

						snake_case_				:					Union[str, Any]   =       SparkExamplesIterable(__magic_name__						).shuffle_data_sources(__magic_name__						)
						assert shuffled_it.n_shards == 3
						for i, (row_id, row_dict) in enumerate(__magic_name__						):
									snake_case_,	snake_case_				:					Tuple   =       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 	__UpperCAmelCase       ( )->    Dict:



			"""simple docstring"""



			snake_case_				:					str   =       pyspark.sql.SparkSession.builder.master("local[*]"						).appName("pyspark"						).getOrCreate()
			snake_case_				:					Optional[Any]   =       spark.range(20						).repartition(4						)

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

			# Partitions 1 and 3
			snake_case_				:					List[str]   =       SparkExamplesIterable(__magic_name__						).shard_data_sources(worker_id=1     ,num_workers=2						)
			assert shard_it_a.n_shards == 2
			snake_case_				:					Union[str, Any]   =       _get_expected_row_ids_and_row_dicts_for_partition_order(__magic_name__     ,[1, 3]						)
			for i, (row_id, row_dict) in enumerate(__magic_name__						):
						snake_case_,	snake_case_				:					int   =       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 	__UpperCAmelCase       ( )->    Dict:



			"""simple docstring"""



			snake_case_				:					List[Any]   =       pyspark.sql.SparkSession.builder.master("local[*]"						).appName("pyspark"						).getOrCreate()
			snake_case_				:					List[str]   =       spark.range(100						).repartition(1						)
			snake_case_				:					Optional[Any]   =       Spark(__magic_name__						)
			# 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












'''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 argparse
import os

from accelerate.test_utils import execute_subprocess_async

def 	__UpperCAmelCase       ( __magic_name__=None						)->    List[str]:



			"""simple docstring"""



			if subparsers is not None:
						snake_case_				:					List[str]   =       subparsers.add_parser("test"						)
			else:
						snake_case_				:					List[Any]   =       argparse.ArgumentParser("Accelerate test command"						)

			parser.add_argument(
			    "--config_file"     ,default=__magic_name__     ,help=(
			        "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
			        "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
			        "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
			        "with 'huggingface'."
			    )     ,)

			if subparsers is not None:
						parser.set_defaults(func=__magic_name__						)
			return parser

def 	__UpperCAmelCase       ( __magic_name__						)->    Tuple:



			"""simple docstring"""



			snake_case_				:					Optional[Any]   =       os.path.sep.join(__file__.split(os.path.sep						)[:-2] + ["test_utils", "scripts", "test_script.py"]						)

			if args.config_file is None:
						snake_case_				:					str   =       script_name
			else:
						snake_case_				:					Any   =       F'''--config_file={args.config_file} {script_name}'''

			snake_case_				:					Union[str, Any]   =       ["accelerate-launch"] + test_args.split()
			snake_case_				:					Optional[int]   =       execute_subprocess_async(__magic_name__     ,env=os.environ.copy()						)
			if result.returncode == 0:
						print("Test is a success! You are ready for your distributed training!"						)

def 	__UpperCAmelCase       ( )->    int:



			"""simple docstring"""



			snake_case_				:					Dict   =       test_command_parser()
			snake_case_				:					Dict   =       parser.parse_args()
			test_command(__magic_name__						)


if __name__ == "__main__":
		main()




