






"""simple docstring"""




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)

_snake_case          =   'bert-base-cased'
_snake_case          =   'fp16'
_snake_case          =   'bf16'
_snake_case          =   [FPaa, BFaa]






@require_fsdp
@require_cuda
class        UpperCamelCase			(	__SCREAMING_SNAKE_CASE      ):


 def _lowercase	(     self		:						List[str]							)					->       Union[str, Any]:
   super().setUp()

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


 def _lowercase	(     self		:						int							)					->       Optional[Any]:
   from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

   for i, strategy in enumerate(_SCREAMING_SNAKE_CASE							):
     _a  :       Tuple             =	self.dist_env.copy()
     _a  :       Dict             =	f"""{i + 1}"""
     _a  :       Union[str, Any]             =	strategy
     with mockenv_context(**_SCREAMING_SNAKE_CASE							):
       _a  :       str             =	FullyShardedDataParallelPlugin()
       self.assertEqual(fsdp_plugin.sharding_strategy							,						ShardingStrategy(i + 1							)							)


 def _lowercase	(     self		:						List[str]							)					->       Union[str, Any]:
   from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

   for i, prefetch_policy in enumerate(_SCREAMING_SNAKE_CASE							):
     _a  :       int             =	self.dist_env.copy()
     _a  :       List[str]             =	prefetch_policy
     with mockenv_context(**_SCREAMING_SNAKE_CASE							):
       _a  :       Dict             =	FullyShardedDataParallelPlugin()
       if prefetch_policy == "NO_PREFETCH":
         self.assertIsNone(fsdp_plugin.backward_prefetch							)
       else:
         self.assertEqual(fsdp_plugin.backward_prefetch							,						BackwardPrefetch(i + 1							)							)


 def _lowercase	(     self		:						int							)					->       Union[str, Any]:
   from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

   for i, state_dict_type in enumerate(_SCREAMING_SNAKE_CASE							):
     _a  :       Union[str, Any]             =	self.dist_env.copy()
     _a  :       Optional[Any]             =	state_dict_type
     with mockenv_context(**_SCREAMING_SNAKE_CASE							):
       _a  :       int             =	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 _lowercase	(     self		:						Dict							)					->       Dict:
   _a  :       List[str]             =	AutoModel.from_pretrained(_SCREAMING_SNAKE_CASE							)
   for policy in FSDP_AUTO_WRAP_POLICY:
     _a  :       List[str]             =	self.dist_env.copy()
     _a  :       List[str]             =	policy
     if policy == "TRANSFORMER_BASED_WRAP":
       _a  :       List[str]             =	"BertLayer"
     elif policy == "SIZE_BASED_WRAP":
       _a  :       Any             =	"2000"
     with mockenv_context(**_SCREAMING_SNAKE_CASE							):
       _a  :       Dict             =	FullyShardedDataParallelPlugin()
       fsdp_plugin.set_auto_wrap_policy(_SCREAMING_SNAKE_CASE							)
       if policy == "NO_WRAP":
         self.assertIsNone(fsdp_plugin.auto_wrap_policy							)
       else:
         self.assertIsNotNone(fsdp_plugin.auto_wrap_policy							)

   _a  :       List[str]             =	self.dist_env.copy()
   _a  :       int             =	"TRANSFORMER_BASED_WRAP"
   _a  :       List[Any]             =	"T5Layer"
   with mockenv_context(**_SCREAMING_SNAKE_CASE							):
     _a  :       Optional[Any]             =	FullyShardedDataParallelPlugin()
     with self.assertRaises(_SCREAMING_SNAKE_CASE							) as cm:
       fsdp_plugin.set_auto_wrap_policy(_SCREAMING_SNAKE_CASE							)
     self.assertTrue("""Could not find the transformer layer class to wrap in the model.""" in str(cm.exception							)							)

   _a  :       Any             =	self.dist_env.copy()
   _a  :       Any             =	"SIZE_BASED_WRAP"
   _a  :       str             =	"0"
   with mockenv_context(**_SCREAMING_SNAKE_CASE							):
     _a  :       Tuple             =	FullyShardedDataParallelPlugin()
     fsdp_plugin.set_auto_wrap_policy(_SCREAMING_SNAKE_CASE							)
     self.assertIsNone(fsdp_plugin.auto_wrap_policy							)


 def _lowercase	(     self		:						Optional[Any]							)					->       Union[str, Any]:
   from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
   from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

   for mp_dtype in dtypes:
     _a  :       Union[str, Any]             =	self.dist_env.copy()
     _a  :       Tuple             =	mp_dtype
     with mockenv_context(**_SCREAMING_SNAKE_CASE							):
       _a  :       Any             =	Accelerator()
       if mp_dtype == "fp16":
         _a  :       Optional[Any]             =	torch.floataa
       elif mp_dtype == "bf16":
         _a  :       List[Any]             =	torch.bfloataa
       _a  :       int             =	MixedPrecision(param_dtype=_SCREAMING_SNAKE_CASE							,						reduce_dtype=_SCREAMING_SNAKE_CASE							,						buffer_dtype=_SCREAMING_SNAKE_CASE							)
       self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy							,						_SCREAMING_SNAKE_CASE							)
       if mp_dtype == FPaa:
         self.assertTrue(isinstance(accelerator.scaler							,						_SCREAMING_SNAKE_CASE							)							)
       elif mp_dtype == BFaa:
         self.assertIsNone(accelerator.scaler							)
       AcceleratorState._reset_state(_SCREAMING_SNAKE_CASE							)







 def _lowercase	(     self		:						Optional[Any]							)					->       int:
   from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

   for flag in [True, False]:
     _a  :       Tuple             =	self.dist_env.copy()
     _a  :       List[Any]             =	str(_SCREAMING_SNAKE_CASE							).lower()
     with mockenv_context(**_SCREAMING_SNAKE_CASE							):
       _a  :       int             =	FullyShardedDataParallelPlugin()
       self.assertEqual(fsdp_plugin.cpu_offload							,						CPUOffload(offload_params=_SCREAMING_SNAKE_CASE							)							)







@require_fsdp
@require_multi_gpu
@slow
class        UpperCamelCase			(	__SCREAMING_SNAKE_CASE      ):


 def _lowercase	(     self		:						List[str]							)					->       List[str]:
   super().setUp()
   _a  :       Optional[Any]             =	0.8_2
   _a  :       Dict             =	[
       "fsdp_shard_grad_op_transformer_based_wrap",
       "fsdp_full_shard_transformer_based_wrap",
   ]
   _a  :       List[Any]             =	{
       "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
   }
   _a  :       Tuple             =	160
   _a  :       Any             =	160

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


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


 def _lowercase	(     self		:						List[Any]							)					->       Dict:
   _a  :       List[str]             =	os.path.join(self.test_scripts_folder							,						"""test_checkpointing.py"""							)
   _a  :       List[str]             =	[
       "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(_SCREAMING_SNAKE_CASE							):
     _a  :       Any             =	cmd.copy()
     cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""							)
     if strategy != "FULL_SHARD":
       continue
     _a  :       Union[str, Any]             =	len(_SCREAMING_SNAKE_CASE							)
     for state_dict_type in FSDP_STATE_DICT_TYPE:
       _a  :       str             =	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(_SCREAMING_SNAKE_CASE							,						env=os.environ.copy()							)

       _a  :       Tuple             =	cmd_config[:-1]
       _a  :       List[str]             =	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(_SCREAMING_SNAKE_CASE							,						env=os.environ.copy()							)







 def _lowercase	(     self		:						List[Any]							)					->       Any:
   _a  :       Union[str, Any]             =	os.path.join(self.test_scripts_folder							,						"""test_peak_memory_usage.py"""							)
   _a  :       str             =	[
       "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():
     _a  :       Dict             =	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(_SCREAMING_SNAKE_CASE							):
         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(_SCREAMING_SNAKE_CASE							,						env=os.environ.copy()							)






"""simple docstring"""






from __future__ import annotations

import time

import numpy as np

__snake_case							:   Optional[Any]      =    [8, 5, 9, 7]
__snake_case							:   List[Any]      =    [
    [2, 0, 1, 1],
    [0, 1, 2, 1],
    [4, 0, 0, 3],
    [0, 2, 1, 0],
    [1, 0, 3, 0],
]
__snake_case							:   Optional[int]      =    [
    [3, 2, 1, 4],
    [0, 2, 5, 2],
    [5, 1, 0, 5],
    [1, 5, 3, 0],
    [3, 0, 3, 3],
]



class 			A__   :




	'''simple docstring'''





	def __init__(     self:		Any       ,			_SCREAMING_SNAKE_CASE:		list[int]       ,			_SCREAMING_SNAKE_CASE:		list[list[int]]       ,			_SCREAMING_SNAKE_CASE:		list[list[int]]       ,			)				->       None:

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       claim_vector
				__lowerCAmelCase      :  Tuple												=       allocated_resources_table
				__lowerCAmelCase      :  Tuple												=       maximum_claim_table

	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       list[int]:

				"""simple docstring"""

				return [
				    sum(p_item[i] for p_item in self.__allocated_resources_table)
				    for i in range(len(self.__allocated_resources_table[0]))
				]

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       list[int]:

				"""simple docstring"""

				return np.array(self.__claim_vector) - np.array(
				    self.__processes_resource_summation())

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       list[list[int]]:

				"""simple docstring"""

				return [
				    list(np.array(self.__maximum_claim_table[i]) - np.array(_SCREAMING_SNAKE_CASE))
				    for i, allocated_resource in enumerate(self.__allocated_resources_table)
				]

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       dict[int, list[int]]:

				"""simple docstring"""

				return {self.__need().index(_SCREAMING_SNAKE_CASE): i for i in self.__need()}

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			**_SCREAMING_SNAKE_CASE:		List[Any])				->       None:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[int]												=       self.__need()
				__lowerCAmelCase      :  int												=       self.__allocated_resources_table
				__lowerCAmelCase      :  Dict												=       self.__available_resources()
				__lowerCAmelCase      :  str												=       self.__need_index_manager()
				for kw, val in kwargs.items():
							if kw and val is True:
										self.__pretty_data()
				print("_" * 50 + "\n")
				while need_list:
							__lowerCAmelCase      :  int												=       False
							for each_need in need_list:
										__lowerCAmelCase      :  Dict												=       True
										for index, need in enumerate(_SCREAMING_SNAKE_CASE):
													if need > available_resources[index]:
																__lowerCAmelCase      :  Dict												=       False
																break
										if execution:
													__lowerCAmelCase      :  Any												=       True
													# get the original index of the process from ind_ctrl db
													for original_need_index, need_clone in need_index_manager.items():
																if each_need == need_clone:
																			__lowerCAmelCase      :  Union[str, Any]												=       original_need_index
													print(F"""Process {process_number + 1} is executing.""")
													# remove the process run from stack
													need_list.remove(_SCREAMING_SNAKE_CASE)
													# update available/freed resources stack
													__lowerCAmelCase      :  Dict												=       np.array(_SCREAMING_SNAKE_CASE) + np.array(
													    alloc_resources_table[process_number])
													print(
													    "Updated available resource stack for processes: "
													    + " ".join([str(_SCREAMING_SNAKE_CASE) for x in available_resources]))
													break
							if safe:
										print("The process is in a safe state.\n")
							else:
										print("System in unsafe state. Aborting...\n")
										break







	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       List[Any]:

				"""simple docstring"""

				print(" " * 9 + "Allocated Resource Table")
				for item in self.__allocated_resources_table:
							print(
							    F"""P{self.__allocated_resources_table.index(_SCREAMING_SNAKE_CASE) + 1}"""
							    + " ".join(F"""{it:>8}""" for it in item)
							    + "\n")
				print(" " * 9 + "System Resource Table")
				for item in self.__maximum_claim_table:
							print(
							    F"""P{self.__maximum_claim_table.index(_SCREAMING_SNAKE_CASE) + 1}"""
							    + " ".join(F"""{it:>8}""" for it in item)
							    + "\n")
				print(
				    "Current Usage by Active Processes: "
				    + " ".join(str(_SCREAMING_SNAKE_CASE) for x in self.__claim_vector))
				print(
				    "Initial Available Resources:       "
				    + " ".join(str(_SCREAMING_SNAKE_CASE) for x in self.__available_resources()))
				time.sleep(1)


if __name__ == "__main__":
	import doctest

	doctest.testmod()