


"""simple docstring"""

import operator as op


UpperCAmelCase__  							="scaler.pt"
UpperCAmelCase__  							="pytorch_model"
UpperCAmelCase__  							="random_states"
UpperCAmelCase__  							="optimizer"
UpperCAmelCase__  							="scheduler"
UpperCAmelCase__  							="pytorch_model.bin"
UpperCAmelCase__  							="pytorch_model.bin.index.json"
UpperCAmelCase__  							="model.safetensors"
UpperCAmelCase__  							="model.safetensors.index.json"
UpperCAmelCase__  							="1.10.2"
UpperCAmelCase__  							="py38"
UpperCAmelCase__  							="4.17.0"
UpperCAmelCase__  							=["ml.p3.16xlarge", "ml.p3dn.24xlarge", "ml.p4dn.24xlarge"]
UpperCAmelCase__  							=["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD", "HYBRID_SHARD", "HYBRID_SHARD_ZERO2"]
UpperCAmelCase__  							=["TRANSFORMER_BASED_WRAP", "SIZE_BASED_WRAP", "NO_WRAP"]
UpperCAmelCase__  							=["BACKWARD_PRE", "BACKWARD_POST", "NO_PREFETCH"]
UpperCAmelCase__  							=["FULL_STATE_DICT", "LOCAL_STATE_DICT", "SHARDED_STATE_DICT"]
UpperCAmelCase__  							="2.0.1"
UpperCAmelCase__  							=["pdsh", "standard", "openmpi", "mvapich"]
UpperCAmelCase__  							=["default", "reduce-overhead", "max-autotune"]

UpperCAmelCase__  							={">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
UpperCAmelCase__  							=[
    "nnodes",
    "nproc_per_node",
    "rdzv_backend",
    "rdzv_endpoint",
    "rdzv_id",
    "rdzv_conf",
    "standalone",
    "max_restarts",
    "monitor_interval",
    "start_method",
    "role",
    "module",
    "m",
    "no_python",
    "run_path",
    "log_dir",
    "r",
    "redirects",
    "t",
    "tee",
    "node_rank",
    "master_addr",
    "master_port",
]

UpperCAmelCase__  							=["DEEPSPEED", "MULTI_GPU", "FSDP", "MEGATRON_LM"]
UpperCAmelCase__  							=["DEEPSPEED", "MULTI_XPU", "FSDP"]




"""simple docstring"""

from typing import List

from .keymap import KEYMAP, get_character


def lowerCAmelCase_      (				UpperCamelCase__	: str ):






 """simple docstring"""







 def decorator(UpperCamelCase__	: Tuple ):
  __lowercase									=	getattr(UpperCamelCase__				,   """handle_key"""				,   [] )
  handle += [key]
  setattr(UpperCamelCase__				,   """handle_key"""				,   UpperCamelCase__ )
  return func

 return decorator


def lowerCAmelCase_      (				*UpperCamelCase__	: List[str] ):






 """simple docstring"""







 def decorator(UpperCamelCase__	: Tuple ):
  __lowercase									=	getattr(UpperCamelCase__				,   """handle_key"""				,   [] )
  handle += keys
  setattr(UpperCamelCase__				,   """handle_key"""				,   UpperCamelCase__ )
  return func

 return decorator





class    lowerCamelCase__						(		_a    ):
      def __new__(		cls	:	str      ,       A_	:	Optional[Any]      ,       A_	:	Union[str, Any]      ,       A_	:	int							):






       '''simple docstring'''


       __lowercase									=	super().__new__(cls      ,       A_      ,       A_      ,       A_							)
       if not hasattr(A_      ,       """key_handler"""							):
        setattr(A_      ,       """key_handler"""      ,       {}							)
       setattr(A_      ,       """handle_input"""      ,       KeyHandler.handle_input							)

       for value in attrs.values():
        __lowercase									=	getattr(A_      ,       """handle_key"""      ,       []							)
        for key in handled_keys:
         __lowercase									=	value
       return new_cls






      @staticmethod
      def SCREAMING_SNAKE_CASE_	(		cls	:	Dict							):






       '''simple docstring'''


       __lowercase									=	get_character()
       if char != KEYMAP["undefined"]:
        __lowercase									=	ord(A_							)
       __lowercase									=	cls.key_handler.get(A_							)
       if handler:
        __lowercase									=	char
        return handler(cls							)
       else:
        return None






def lowerCAmelCase_      (				cls	: int ):






 """simple docstring"""







 return KeyHandler(cls.__name__				,   cls.__bases__				,   cls.__dict__.copy() )
