import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(42)

__lowerCamelCase    :				Union[str, Any]							     =	'''bert-base-cased'''
__lowerCamelCase    :				Dict							     =	'''fp16'''
__lowerCamelCase    :				Dict							     =	'''bf16'''
__lowerCamelCase    :				List[Any]							     =	[FPaa, BFaa]







@require_fsdp
@require_cuda
class __snake_case       (		lowerCamelCase_	):




   def 			__a		(					self    :       Union[str, Any]   ):



     """simple docstring"""






     super().setUp()

     SCREAMING_SNAKE_CASE__    							=   dict(
         ACCELERATE_USE_FSDP="""true""" ,			MASTER_ADDR="""localhost""" ,			MASTER_PORT="""10999""" ,			RANK="""0""" ,			LOCAL_RANK="""0""" ,			WORLD_SIZE="""1""" ,			)




   def 			__a		(					self    :       Optional[int]   ):



     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

     for i, strategy in enumerate(_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
       SCREAMING_SNAKE_CASE__    							=   f"""{i + 1}"""
       SCREAMING_SNAKE_CASE__    							=   strategy
       with mockenv_context(**_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   FullyShardedDataParallelPlugin()
         self.assertEqual(fsdp_plugin.sharding_strategy ,			ShardingStrategy(i + 1   )   )




   def 			__a		(					self    :       str   ):



     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

     for i, prefetch_policy in enumerate(_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
       SCREAMING_SNAKE_CASE__    							=   prefetch_policy
       with mockenv_context(**_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   FullyShardedDataParallelPlugin()
         if prefetch_policy == "NO_PREFETCH":
           self.assertIsNone(fsdp_plugin.backward_prefetch   )
         else:
           self.assertEqual(fsdp_plugin.backward_prefetch ,			BackwardPrefetch(i + 1   )   )




   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

     for i, state_dict_type in enumerate(_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
       SCREAMING_SNAKE_CASE__    							=   state_dict_type
       with mockenv_context(**_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   FullyShardedDataParallelPlugin()
         self.assertEqual(fsdp_plugin.state_dict_type ,			StateDictType(i + 1   )   )
         if state_dict_type == "FULL_STATE_DICT":
           self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu   )
           self.assertTrue(fsdp_plugin.state_dict_config.ranka_only   )




   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   AutoModel.from_pretrained(_lowercase   )
     for policy in FSDP_AUTO_WRAP_POLICY:
       SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
       SCREAMING_SNAKE_CASE__    							=   policy
       if policy == "TRANSFORMER_BASED_WRAP":
         SCREAMING_SNAKE_CASE__    							=   """BertLayer"""
       elif policy == "SIZE_BASED_WRAP":
         SCREAMING_SNAKE_CASE__    							=   """2000"""
       with mockenv_context(**_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   FullyShardedDataParallelPlugin()
         fsdp_plugin.set_auto_wrap_policy(_lowercase   )
         if policy == "NO_WRAP":
           self.assertIsNone(fsdp_plugin.auto_wrap_policy   )
         else:
           self.assertIsNotNone(fsdp_plugin.auto_wrap_policy   )

     SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
     SCREAMING_SNAKE_CASE__    							=   """TRANSFORMER_BASED_WRAP"""
     SCREAMING_SNAKE_CASE__    							=   """T5Layer"""
     with mockenv_context(**_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   FullyShardedDataParallelPlugin()
       with self.assertRaises(_lowercase   ) as cm:
         fsdp_plugin.set_auto_wrap_policy(_lowercase   )
       self.assertTrue("""Could not find the transformer layer class to wrap in the model.""" in str(cm.exception   )   )

     SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
     SCREAMING_SNAKE_CASE__    							=   """SIZE_BASED_WRAP"""
     SCREAMING_SNAKE_CASE__    							=   """0"""
     with mockenv_context(**_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   FullyShardedDataParallelPlugin()
       fsdp_plugin.set_auto_wrap_policy(_lowercase   )
       self.assertIsNone(fsdp_plugin.auto_wrap_policy   )




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
     from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

     for mp_dtype in dtypes:
       SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
       SCREAMING_SNAKE_CASE__    							=   mp_dtype
       with mockenv_context(**_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   Accelerator()
         if mp_dtype == "fp16":
           SCREAMING_SNAKE_CASE__    							=   torch.floataa
         elif mp_dtype == "bf16":
           SCREAMING_SNAKE_CASE__    							=   torch.bfloataa
         SCREAMING_SNAKE_CASE__    							=   MixedPrecision(param_dtype=_lowercase ,			reduce_dtype=_lowercase ,			buffer_dtype=_lowercase   )
         self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy ,			_lowercase   )
         if mp_dtype == FPaa:
           self.assertTrue(isinstance(accelerator.scaler ,			_lowercase   )   )
         elif mp_dtype == BFaa:
           self.assertIsNone(accelerator.scaler   )
         AcceleratorState._reset_state(_lowercase   )




   def 			__a		(					self    :       str   ):



     """simple docstring"""






     from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

     for flag in [True, False]:
       SCREAMING_SNAKE_CASE__    							=   self.dist_env.copy()
       SCREAMING_SNAKE_CASE__    							=   str(_lowercase   ).lower()
       with mockenv_context(**_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   FullyShardedDataParallelPlugin()
         self.assertEqual(fsdp_plugin.cpu_offload ,			CPUOffload(offload_params=_lowercase   )   )







@require_fsdp
@require_multi_gpu
@slow
class __snake_case       (		lowerCamelCase_	):




   def 			__a		(					self    :       Union[str, Any]   ):



     """simple docstring"""






     super().setUp()
     SCREAMING_SNAKE_CASE__    							=   0.82
     SCREAMING_SNAKE_CASE__    							=   [
         """fsdp_shard_grad_op_transformer_based_wrap""",
         """fsdp_full_shard_transformer_based_wrap""",
     ]
     SCREAMING_SNAKE_CASE__    							=   {
         """multi_gpu_fp16""": 32_00,
         """fsdp_shard_grad_op_transformer_based_wrap_fp16""": 20_00,
         """fsdp_full_shard_transformer_based_wrap_fp16""": 19_00,
         # Disabling below test as it overwhelms the RAM memory usage
         # on CI self-hosted runner leading to tests getting killed.
         # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
     }
     SCREAMING_SNAKE_CASE__    							=   1_60
     SCREAMING_SNAKE_CASE__    							=   1_60

     SCREAMING_SNAKE_CASE__    							=   inspect.getfile(accelerate.test_utils   )
     SCREAMING_SNAKE_CASE__    							=   os.path.sep.join(mod_file.split(os.path.sep   )[:-1] + ["""scripts""", """external_deps"""]   )




   def 			__a		(					self    :       Union[str, Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   os.path.join(self.test_scripts_folder ,			"""test_performance.py"""   )
     SCREAMING_SNAKE_CASE__    							=   ["""accelerate""", """launch""", """--num_processes=2""", """--num_machines=1""", """--machine_rank=0""", """--use_fsdp"""]
     for config in self.performance_configs:
       SCREAMING_SNAKE_CASE__    							=   cmd.copy()
       for i, strategy in enumerate(_lowercase   ):
         if strategy.lower() in config:
           cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""   )
           break

       if "fp32" in config:
         cmd_config.append("""--mixed_precision=no"""   )
       else:
         cmd_config.append("""--mixed_precision=fp16"""   )

       if "cpu_offload" in config:
         cmd_config.append("""--fsdp_offload_params=True"""   )

       for policy in FSDP_AUTO_WRAP_POLICY:
         if policy.lower() in config:
           cmd_config.append(f"""--fsdp_auto_wrap_policy={policy}"""   )
           break

       if policy == "TRANSFORMER_BASED_WRAP":
         cmd_config.append("""--fsdp_transformer_layer_cls_to_wrap=BertLayer"""   )
       elif policy == "SIZE_BASED_WRAP":
         cmd_config.append("""--fsdp_min_num_params=2000"""   )

       cmd_config.extend(
           [
               self.test_file_path,
               f"""--output_dir={self.tmpdir}""",
               f"""--performance_lower_bound={self.performance_lower_bound}""",
           ]   )
       with patch_environment(omp_num_threads=1   ):
         execute_subprocess_async(_lowercase ,			env=os.environ.copy()   )




   def 			__a		(					self    :       str   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   os.path.join(self.test_scripts_folder ,			"""test_checkpointing.py"""   )
     SCREAMING_SNAKE_CASE__    							=   [
         """accelerate""",
         """launch""",
         """--num_processes=2""",
         """--num_machines=1""",
         """--machine_rank=0""",
         """--use_fsdp""",
         """--mixed_precision=fp16""",
         """--fsdp_transformer_layer_cls_to_wrap=BertLayer""",
     ]

     for i, strategy in enumerate(_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   cmd.copy()
       cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""   )
       if strategy != "FULL_SHARD":
         continue
       SCREAMING_SNAKE_CASE__    							=   len(_lowercase   )
       for state_dict_type in FSDP_STATE_DICT_TYPE:
         SCREAMING_SNAKE_CASE__    							=   cmd_config[:state_dict_config_index]
         cmd_config.append(f"""--fsdp_state_dict_type={state_dict_type}"""   )
         cmd_config.extend(
             [
                 self.test_file_path,
                 f"""--output_dir={self.tmpdir}""",
                 """--partial_train_epoch=1""",
             ]   )
         with patch_environment(omp_num_threads=1   ):
           execute_subprocess_async(_lowercase ,			env=os.environ.copy()   )

         SCREAMING_SNAKE_CASE__    							=   cmd_config[:-1]
         SCREAMING_SNAKE_CASE__    							=   os.path.join(self.tmpdir ,			"""epoch_0"""   )
         cmd_config.extend(
             [
                 f"""--resume_from_checkpoint={resume_from_checkpoint}""",
             ]   )
         with patch_environment(omp_num_threads=1   ):
           execute_subprocess_async(_lowercase ,			env=os.environ.copy()   )




   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   os.path.join(self.test_scripts_folder ,			"""test_peak_memory_usage.py"""   )
     SCREAMING_SNAKE_CASE__    							=   [
         """accelerate""",
         """launch""",
         """--num_processes=2""",
         """--num_machines=1""",
         """--machine_rank=0""",
     ]
     for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
       SCREAMING_SNAKE_CASE__    							=   cmd.copy()
       if "fp16" in spec:
         cmd_config.extend(["""--mixed_precision=fp16"""]   )
       else:
         cmd_config.extend(["""--mixed_precision=no"""]   )

       if "multi_gpu" in spec:
         continue
       else:
         cmd_config.extend(["""--use_fsdp"""]   )
         for i, strategy in enumerate(_lowercase   ):
           if strategy.lower() in spec:
             cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""   )
             break

         if "cpu_offload" in spec:
           cmd_config.append("""--fsdp_offload_params=True"""   )

         for policy in FSDP_AUTO_WRAP_POLICY:
           if policy.lower() in spec:
             cmd_config.append(f"""--fsdp_auto_wrap_policy={policy}"""   )
             break

         if policy == "TRANSFORMER_BASED_WRAP":
           cmd_config.append("""--fsdp_transformer_layer_cls_to_wrap=BertLayer"""   )
         elif policy == "SIZE_BASED_WRAP":
           cmd_config.append("""--fsdp_min_num_params=2000"""   )

       cmd_config.extend(
           [
               self.test_file_path,
               f"""--output_dir={self.tmpdir}""",
               f"""--peak_memory_upper_bound={peak_mem_upper_bound}""",
               f"""--n_train={self.n_train}""",
               f"""--n_val={self.n_val}""",
           ]   )
       with patch_environment(omp_num_threads=1   ):
         execute_subprocess_async(_lowercase ,			env=os.environ.copy()   )

import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


__lowerCamelCase    :				Union[str, Any]							     =	datasets.utils.logging.get_logger(__name__)







@dataclass
class __snake_case       (		datasets.BuilderConfig	):
   lowerCAmelCase_														=							1_00_00
   lowerCAmelCase_														=							None
   lowerCAmelCase_														=							None







class __snake_case       (		datasets.ArrowBasedBuilder	):
   lowerCAmelCase_														=							ParquetConfig




   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     return datasets.DatasetInfo(features=self.config.features   )




   def 			__a		(					self    :       Dict ,			_lowercase    :       List[str]   ):



     """simple docstring"""






     if not self.config.data_files:
       raise ValueError(f"""At least one data file must be specified, but got data_files={self.config.data_files}"""   )
     SCREAMING_SNAKE_CASE__    							=   dl_manager.download_and_extract(self.config.data_files   )
     if isinstance(_lowercase ,			(str, list, tuple)   ):
       SCREAMING_SNAKE_CASE__    							=   data_files
       if isinstance(_lowercase ,			_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   [files]
       # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
       SCREAMING_SNAKE_CASE__    							=   [dl_manager.iter_files(_lowercase   ) for file in files]
       return [datasets.SplitGenerator(name=datasets.Split.TRAIN ,			gen_kwargs={"""files""": files}   )]
     SCREAMING_SNAKE_CASE__    							=   []
     for split_name, files in data_files.items():
       if isinstance(_lowercase ,			_lowercase   ):
         SCREAMING_SNAKE_CASE__    							=   [files]
       # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
       SCREAMING_SNAKE_CASE__    							=   [dl_manager.iter_files(_lowercase   ) for file in files]
       # Infer features is they are stoed in the arrow schema
       if self.info.features is None:
         for file in itertools.chain.from_iterable(_lowercase   ):
           with open(_lowercase ,			"""rb"""   ) as f:
             SCREAMING_SNAKE_CASE__    							=   datasets.Features.from_arrow_schema(pq.read_schema(_lowercase   )   )
           break
       splits.append(datasets.SplitGenerator(name=_lowercase ,			gen_kwargs={"""files""": files}   )   )
     return splits




   def 			__a		(					self    :       Optional[Any] ,			_lowercase    :       pa.Table   ):



     """simple docstring"""






     if self.info.features is not None:
       # more expensive cast to support nested features with keys in a different order
       # allows str <-> int/float or str to Audio for example
       SCREAMING_SNAKE_CASE__    							=   table_cast(_lowercase ,			self.info.features.arrow_schema   )
     return pa_table




   def 			__a		(					self    :       Dict ,			_lowercase    :       List[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.info.features.arrow_schema if self.info.features is not None else None
     if self.info.features is not None and self.config.columns is not None:
       if sorted(field.name for field in schema   ) != sorted(self.config.columns   ):
         raise ValueError(
             f"""Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'"""   )
     for file_idx, file in enumerate(itertools.chain.from_iterable(_lowercase   )   ):
       with open(_lowercase ,			"""rb"""   ) as f:
         SCREAMING_SNAKE_CASE__    							=   pq.ParquetFile(_lowercase   )
         try:
           for batch_idx, record_batch in enumerate(
               parquet_file.iter_batches(batch_size=self.config.batch_size ,			columns=self.config.columns   )   ):
             SCREAMING_SNAKE_CASE__    							=   pa.Table.from_batches([record_batch]   )
             # Uncomment for debugging (will print the Arrow table size and elements)
             # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
             # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
             yield f"""{file_idx}_{batch_idx}""", self._cast_table(_lowercase   )
         except ValueError as e:
           logger.error(f"""Failed to read file '{file}' with error {type(_lowercase   )}: {e}"""   )
           raise
