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(4_2)

_SCREAMING_SNAKE_CASE      						=	"bert-base-cased"
_SCREAMING_SNAKE_CASE      						=	"fp16"
_SCREAMING_SNAKE_CASE      						=	"bf16"
_SCREAMING_SNAKE_CASE      						=	[FPaa, BFaa]
@require_fsdp
@require_cuda
class SCREAMING_SNAKE_CASE_						(	lowerCAmelCase__						):







   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    super().setUp()

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







   def 		lowerCamelCase_     (			self		:		Optional[Any]							):





    """simple docstring"""


    from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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







   def 		lowerCamelCase_     (			self		:		Optional[int]							):





    """simple docstring"""


    from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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







   def 		lowerCamelCase_     (			self		:		Optional[int]							):





    """simple docstring"""


    from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

    for i, state_dict_type in enumerate(__a							):
     UpperCamelCase          =	self.dist_env.copy()
     UpperCamelCase          =	state_dict_type
     with mockenv_context(**__a							):
      UpperCamelCase          =	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 		lowerCamelCase_     (			self		:		Optional[Any]							):





    """simple docstring"""


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

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

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







   def 		lowerCamelCase_     (			self		:		Any							):





    """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:
     UpperCamelCase          =	self.dist_env.copy()
     UpperCamelCase          =	mp_dtype
     with mockenv_context(**__a							):
      UpperCamelCase          =	Accelerator()
      if mp_dtype == "fp16":
       UpperCamelCase          =	torch.floataa
      elif mp_dtype == "bf16":
       UpperCamelCase          =	torch.bfloataa
      UpperCamelCase          =	MixedPrecision(param_dtype=__a					,					reduce_dtype=__a					,					buffer_dtype=__a							)
      self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy					,					__a							)
      if mp_dtype == FPaa:
       self.assertTrue(isinstance(accelerator.scaler					,					__a							)							)
      elif mp_dtype == BFaa:
       self.assertIsNone(accelerator.scaler							)
      AcceleratorState._reset_state(__a							)







   def 		lowerCamelCase_     (			self		:		Tuple							):





    """simple docstring"""


    from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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





@require_fsdp
@require_multi_gpu
@slow
class SCREAMING_SNAKE_CASE_						(	lowerCAmelCase__						):







   def 		lowerCamelCase_     (			self		:		int							):





    """simple docstring"""


    super().setUp()
    UpperCamelCase          =	0.8_2
    UpperCamelCase          =	[
        """fsdp_shard_grad_op_transformer_based_wrap""",
        """fsdp_full_shard_transformer_based_wrap""",
    ]
    UpperCamelCase          =	{
        """multi_gpu_fp16""": 3200,
        """fsdp_shard_grad_op_transformer_based_wrap_fp16""": 2000,
        """fsdp_full_shard_transformer_based_wrap_fp16""": 1900,
        # 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
    }
    UpperCamelCase          =	160
    UpperCamelCase          =	160

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







   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    UpperCamelCase          =	os.path.join(self.test_scripts_folder					,					"""test_performance.py"""							)
    UpperCamelCase          =	["""accelerate""", """launch""", """--num_processes=2""", """--num_machines=1""", """--machine_rank=0""", """--use_fsdp"""]
    for config in self.performance_configs:
     UpperCamelCase          =	cmd.copy()
     for i, strategy in enumerate(__a							):
      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(__a					,					env=os.environ.copy()							)







   def 		lowerCamelCase_     (			self		:		List[str]							):





    """simple docstring"""


    UpperCamelCase          =	os.path.join(self.test_scripts_folder					,					"""test_checkpointing.py"""							)
    UpperCamelCase          =	[
        """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(__a							):
     UpperCamelCase          =	cmd.copy()
     cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""							)
     if strategy != "FULL_SHARD":
      continue
     UpperCamelCase          =	len(__a							)
     for state_dict_type in FSDP_STATE_DICT_TYPE:
      UpperCamelCase          =	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(__a					,					env=os.environ.copy()							)

      UpperCamelCase          =	cmd_config[:-1]
      UpperCamelCase          =	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(__a					,					env=os.environ.copy()							)







   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    UpperCamelCase          =	os.path.join(self.test_scripts_folder					,					"""test_peak_memory_usage.py"""							)
    UpperCamelCase          =	[
        """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():
     UpperCamelCase          =	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(__a							):
       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(__a					,					env=os.environ.copy()							)




"""simple docstring"""


import argparse
import os
import re

import packaging.version


A     :   Any			    =			"examples/"
A     :   Optional[Any]			    =			{
    "examples": (re.compile(R"^check_min_version\(\"[^\"]+\"\)\s*$", re.MULTILINE), "check_min_version(\"VERSION\")\n"),
    "init": (re.compile(R"^__version__\s+=\s+\"([^\"]+)\"\s*$", re.MULTILINE), "__version__ = \"VERSION\"\n"),
    "setup": (re.compile(R"^(\s*)version\s*=\s*\"[^\"]+\",", re.MULTILINE), R"\1version=\"VERSION\","),
    "doc": (re.compile(R"^(\s*)release\s*=\s*\"[^\"]+\"$", re.MULTILINE), "release = \"VERSION\"\n"),
}
A     :   Optional[int]			    =			{
    "init": "src/transformers/__init__.py",
    "setup": "setup.py",
}
A     :   List[Any]			    =			"README.md"


def _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase    , _UpperCamelCase      ):

			'''simple docstring'''




			with open(_UpperCamelCase    , "r"    , encoding="utf-8"    , newline="\n"      ) as f:
						__lowerCAmelCase				    =			f.read()
			__lowerCAmelCase    ,	__lowerCAmelCase				    =			REPLACE_PATTERNS[pattern]
			__lowerCAmelCase				    =			replace.replace("VERSION"    , _UpperCamelCase      )
			__lowerCAmelCase				    =			re_pattern.sub(_UpperCamelCase    , _UpperCamelCase      )
			with open(_UpperCamelCase    , "w"    , encoding="utf-8"    , newline="\n"      ) as f:
						f.write(_UpperCamelCase      )


def _lowerCamelCase	(		_UpperCamelCase      ):

			'''simple docstring'''




			for folder, directories, fnames in os.walk(_UpperCamelCase      ):
						# Removing some of the folders with non-actively maintained examples from the walk
						if "research_projects" in directories:
									directories.remove("research_projects"      )
						if "legacy" in directories:
									directories.remove("legacy"      )
						for fname in fnames:
									if fname.endswith(".py"      ):
												update_version_in_file(os.path.join(_UpperCamelCase    , _UpperCamelCase      )    , _UpperCamelCase    , pattern="examples"      )


def _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase=False      ):

			'''simple docstring'''




			for pattern, fname in REPLACE_FILES.items():
						update_version_in_file(_UpperCamelCase    , _UpperCamelCase    , _UpperCamelCase      )
			if not patch:
						update_version_in_examples(_UpperCamelCase      )


def _lowerCamelCase	(		):

			'''simple docstring'''




			__lowerCAmelCase				    =			"🤗 Transformers currently provides the following architectures"
			__lowerCAmelCase				    =			"1. Want to contribute a new model?"
			with open(_UpperCamelCase    , "r"    , encoding="utf-8"    , newline="\n"      ) as f:
						__lowerCAmelCase				    =			f.readlines()

			# Find the start of the list.
			__lowerCAmelCase				    =			0
			while not lines[start_index].startswith(_start_prompt      ):
						start_index += 1
			start_index += 1

			__lowerCAmelCase				    =			start_index
			# Update the lines in the model list.
			while not lines[index].startswith(_end_prompt      ):
						if lines[index].startswith("1."      ):
									__lowerCAmelCase				    =			lines[index].replace(
									    "https://huggingface.co/docs/transformers/main/model_doc"    , "https://huggingface.co/docs/transformers/model_doc"    , )
						index += 1

			with open(_UpperCamelCase    , "w"    , encoding="utf-8"    , newline="\n"      ) as f:
						f.writelines(_UpperCamelCase      )


def _lowerCamelCase	(		):

			'''simple docstring'''




			with open(REPLACE_FILES["init"]    , "r"      ) as f:
						__lowerCAmelCase				    =			f.read()
			__lowerCAmelCase				    =			REPLACE_PATTERNS["init"][0].search(_UpperCamelCase      ).groups()[0]
			return packaging.version.parse(_UpperCamelCase      )


def _lowerCamelCase	(		_UpperCamelCase=False      ):

			'''simple docstring'''




			__lowerCAmelCase				    =			get_version()
			if patch and default_version.is_devrelease:
						raise ValueError("Can't create a patch version from the dev branch, checkout a released version!"      )
			if default_version.is_devrelease:
						__lowerCAmelCase				    =			default_version.base_version
			elif patch:
						__lowerCAmelCase				    =			f"{default_version.major}.{default_version.minor}.{default_version.micro + 1}"
			else:
						__lowerCAmelCase				    =			f"{default_version.major}.{default_version.minor + 1}.0"

			# Now let's ask nicely if that's the right one.
			__lowerCAmelCase				    =			input(f"Which version are you releasing? [{default_version}]"      )
			if len(_UpperCamelCase      ) == 0:
						__lowerCAmelCase				    =			default_version

			print(f"Updating version to {version}."      )
			global_version_update(_UpperCamelCase    , patch=_UpperCamelCase      )
			if not patch:
						print("Cleaning main README, don't forget to run `make fix-copies`."      )
						clean_main_ref_in_model_list()


def _lowerCamelCase	(		):

			'''simple docstring'''




			__lowerCAmelCase				    =			get_version()
			__lowerCAmelCase				    =			f"{current_version.major}.{current_version.minor + 1}.0.dev0"
			__lowerCAmelCase				    =			current_version.base_version

			# Check with the user we got that right.
			__lowerCAmelCase				    =			input(f"Which version are we developing now? [{dev_version}]"      )
			if len(_UpperCamelCase      ) == 0:
						__lowerCAmelCase				    =			dev_version

			print(f"Updating version to {version}."      )
			global_version_update(_UpperCamelCase      )
			print("Cleaning main README, don't forget to run `make fix-copies`."      )
			clean_main_ref_in_model_list()


if __name__ == "__main__":
	A     :   Union[str, Any]			    =			argparse.ArgumentParser()
	parser.add_argument("--post_release", action="store_true", help="Whether this is pre or post release.")
	parser.add_argument("--patch", action="store_true", help="Whether or not this is a patch release.")
	A     :   Dict			    =			parser.parse_args()
	if not args.post_release:
		pre_release_work(patch=args.patch)
	elif args.patch:
		print("Nothing to do after a patch :-)")
	else:
		post_release_work()




