
"""simple docstring"""




import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


A__ :       Any				= logging.getLogger(__name__)




def 			_lowerCAmelCase		(		_UpperCamelCase=2			,					_UpperCamelCase=3			,					_UpperCamelCase=16			,					_UpperCamelCase = 10			,					_UpperCamelCase = 2    ):






  """simple docstring"""



  def get_dataset(_UpperCamelCase    ):
    _lowercase:      Tuple   			=  torch.randn(batch_size * n_batches			,					1    )
    return TensorDataset(_UpperCamelCase			,					a * x + b + 0.1 * torch.randn(batch_size * n_batches			,					1    )    )

  _lowercase:      Tuple   			=  get_dataset(_UpperCamelCase    )
  _lowercase:      Dict   			=  get_dataset(_UpperCamelCase    )
  _lowercase:      Optional[int]   			=  DataLoader(_UpperCamelCase			,					shuffle=_UpperCamelCase			,					batch_size=_UpperCamelCase			,					num_workers=4    )
  _lowercase:      Optional[int]   			=  DataLoader(_UpperCamelCase			,					shuffle=_UpperCamelCase			,					batch_size=_UpperCamelCase			,					num_workers=4    )
  return (train_dataloader, valid_dataloader)




def 			_lowerCAmelCase		(		_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase=None    ):






  """simple docstring"""



  _lowercase:      Tuple   			=  []
  for epoch in range(_UpperCamelCase    ):
    # Train quickly
    model.train()
    for batch in dataloader:
      _lowercase			,					_lowercase:      Optional[Any]   			=  batch
      _lowercase:      Union[str, Any]   			=  model(_UpperCamelCase    )
      _lowercase:      Optional[Any]   			=  torch.nn.functional.mse_loss(_UpperCamelCase			,					_UpperCamelCase    )
      accelerator.backward(_UpperCamelCase    )
      optimizer.step()
      optimizer.zero_grad()
    rands.append(random.random()    )  # Introduce some randomness
    if scheduler is not None:
      scheduler.step()
  return rands


class       __magic_name__							(							nn.Module			):
       def __init__(  self							)  ->						str:





         """simple docstring"""




         super().__init__()
         _lowercase:      Optional[Any]   			=  nn.Parameter(torch.randn(1							)							)
         _lowercase:      int   			=  nn.Parameter(torch.randn(1							)							)



       def 				lowercase_				(  self			,					A_							)  ->						List[str]:





         """simple docstring"""




         return x * self.a + self.b


class       __magic_name__							(							unittest.TestCase			):
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         with tempfile.TemporaryDirectory() as tmpdir:
           set_seed(42							)
           _lowercase:      Union[str, Any]   			=  DummyModel()
           _lowercase:      Optional[int]   			=  torch.optim.Adam(params=model.parameters()			,					lr=1E-3							)
           _lowercase			,					_lowercase:      Union[str, Any]   			=  dummy_dataloaders()
           _lowercase:      Any   			=  ProjectConfiguration(total_limit=1			,					project_dir=A_			,					automatic_checkpoint_naming=A_							)
           # Train baseline
           _lowercase:      Dict   			=  Accelerator(project_config=A_							)
           _lowercase			,					_lowercase			,					_lowercase			,					_lowercase:      Optional[int]   			=  accelerator.prepare(
               A_			,					A_			,					A_			,					A_							)
           # Save initial
           accelerator.save_state()

           # Save second state
           accelerator.save_state()
           self.assertEqual(len(os.listdir(accelerator.project_dir							)							)			,					1							)
       def 				lowercase_				(  self							)  ->						int:





         """simple docstring"""




         with tempfile.TemporaryDirectory() as tmpdir:
           set_seed(42							)
           _lowercase:      Dict   			=  DummyModel()
           _lowercase:      str   			=  torch.optim.Adam(params=model.parameters()			,					lr=1E-3							)
           _lowercase			,					_lowercase:      Optional[Any]   			=  dummy_dataloaders()
           # Train baseline
           _lowercase:      Dict   			=  Accelerator()
           _lowercase			,					_lowercase			,					_lowercase			,					_lowercase:      Optional[Any]   			=  accelerator.prepare(
               A_			,					A_			,					A_			,					A_							)
           # Save initial
           _lowercase:      Dict   			=  os.path.join(A_			,					'''initial'''							)
           accelerator.save_state(A_							)
           ((_lowercase)			,					(_lowercase)):      Union[str, Any]   			=  model.a.item(), model.b.item()
           _lowercase:      int   			=  optimizer.state_dict()
           _lowercase:      Optional[int]   			=  train(3			,					A_			,					A_			,					A_			,					A_							)
           ((_lowercase)			,					(_lowercase)):      int   			=  model.a.item(), model.b.item()
           _lowercase:      Optional[int]   			=  optimizer.state_dict()

           # Train partially
           set_seed(42							)
           _lowercase:      Dict   			=  DummyModel()
           _lowercase:      int   			=  torch.optim.Adam(params=model.parameters()			,					lr=1E-3							)
           _lowercase			,					_lowercase:      List[Any]   			=  dummy_dataloaders()
           _lowercase:      str   			=  Accelerator()
           _lowercase			,					_lowercase			,					_lowercase			,					_lowercase:      str   			=  accelerator.prepare(
               A_			,					A_			,					A_			,					A_							)
           accelerator.load_state(A_							)
           ((_lowercase)			,					(_lowercase)):      List[Any]   			=  model.a.item(), model.b.item()
           _lowercase:      Dict   			=  optimizer.state_dict()
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)

           _lowercase:      Optional[int]   			=  train(2			,					A_			,					A_			,					A_			,					A_							)
           # Save everything
           _lowercase:      List[Any]   			=  os.path.join(A_			,					'''checkpoint'''							)
           accelerator.save_state(A_							)

           # Load everything back in and make sure all states work
           accelerator.load_state(A_							)
           test_rands += train(1			,					A_			,					A_			,					A_			,					A_							)
           ((_lowercase)			,					(_lowercase)):      Optional[int]   			=  model.a.item(), model.b.item()
           _lowercase:      Any   			=  optimizer.state_dict()
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         with tempfile.TemporaryDirectory() as tmpdir:
           set_seed(42							)
           _lowercase:      Tuple   			=  DummyModel()
           _lowercase:      int   			=  torch.optim.Adam(params=model.parameters()			,					lr=1E-3							)
           _lowercase			,					_lowercase:      Optional[int]   			=  dummy_dataloaders()
           _lowercase:      Optional[Any]   			=  ProjectConfiguration(automatic_checkpoint_naming=A_							)

           # Train baseline
           _lowercase:      List[Any]   			=  Accelerator(project_dir=A_			,					project_config=A_							)
           _lowercase			,					_lowercase			,					_lowercase			,					_lowercase:      str   			=  accelerator.prepare(
               A_			,					A_			,					A_			,					A_							)
           # Save initial
           accelerator.save_state()
           ((_lowercase)			,					(_lowercase)):      Dict   			=  model.a.item(), model.b.item()
           _lowercase:      Optional[int]   			=  optimizer.state_dict()
           _lowercase:      Tuple   			=  train(3			,					A_			,					A_			,					A_			,					A_							)
           ((_lowercase)			,					(_lowercase)):      int   			=  model.a.item(), model.b.item()
           _lowercase:      Union[str, Any]   			=  optimizer.state_dict()

           # Train partially
           set_seed(42							)
           _lowercase:      Any   			=  DummyModel()
           _lowercase:      int   			=  torch.optim.Adam(params=model.parameters()			,					lr=1E-3							)
           _lowercase			,					_lowercase:      int   			=  dummy_dataloaders()
           _lowercase:      List[Any]   			=  ProjectConfiguration(iteration=1			,					automatic_checkpoint_naming=A_							)
           _lowercase:      Any   			=  Accelerator(project_dir=A_			,					project_config=A_							)
           _lowercase			,					_lowercase			,					_lowercase			,					_lowercase:      str   			=  accelerator.prepare(
               A_			,					A_			,					A_			,					A_							)
           accelerator.load_state(os.path.join(A_			,					'''checkpoints'''			,					'''checkpoint_0'''							)							)
           ((_lowercase)			,					(_lowercase)):      Dict   			=  model.a.item(), model.b.item()
           _lowercase:      Dict   			=  optimizer.state_dict()
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)

           _lowercase:      Any   			=  train(2			,					A_			,					A_			,					A_			,					A_							)
           # Save everything
           accelerator.save_state()

           # Load everything back in and make sure all states work
           accelerator.load_state(os.path.join(A_			,					'''checkpoints'''			,					'''checkpoint_1'''							)							)
           test_rands += train(1			,					A_			,					A_			,					A_			,					A_							)
           ((_lowercase)			,					(_lowercase)):      Optional[int]   			=  model.a.item(), model.b.item()
           _lowercase:      Any   			=  optimizer.state_dict()
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
           self.assertEqual(A_			,					A_							)
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         _lowercase:      Optional[Any]   			=  torch.tensor([1, 2, 3]							)
         _lowercase:      Optional[Any]   			=  torch.tensor([2, 3, 4]							)
         _lowercase:      Union[str, Any]   			=  DummyModel()
         _lowercase:      List[Any]   			=  torch.optim.Adam(net.parameters()							)
         _lowercase:      List[Any]   			=  Accelerator()
         with self.assertRaises(A_							) as ve:
           accelerator.register_for_checkpointing(A_			,					A_			,					A_			,					A_							)
         _lowercase:      Tuple   			=  str(ve.exception							)
         self.assertTrue('''Item at index 0''' in message							)
         self.assertTrue('''Item at index 1''' in message							)
         self.assertFalse('''Item at index 2''' in message							)
         self.assertFalse('''Item at index 3''' in message							)
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         with tempfile.TemporaryDirectory() as tmpdir:
           set_seed(42							)
           _lowercase:      Any   			=  DummyModel()
           _lowercase:      List[Any]   			=  torch.optim.Adam(params=model.parameters()			,					lr=1E-3							)
           _lowercase:      Optional[Any]   			=  torch.optim.lr_scheduler.StepLR(A_			,					step_size=1			,					gamma=0.99							)
           _lowercase			,					_lowercase:      Optional[int]   			=  dummy_dataloaders()
           _lowercase:      Optional[int]   			=  ProjectConfiguration(automatic_checkpoint_naming=A_							)
           # Train baseline
           _lowercase:      int   			=  Accelerator(project_dir=A_			,					project_config=A_							)
           _lowercase			,					_lowercase			,					_lowercase			,					_lowercase			,					_lowercase:      Tuple   			=  accelerator.prepare(
               A_			,					A_			,					A_			,					A_			,					A_							)
           # Save initial
           accelerator.save_state()
           _lowercase:      Union[str, Any]   			=  scheduler.state_dict()
           train(3			,					A_			,					A_			,					A_			,					A_			,					A_							)
           self.assertNotEqual(A_			,					scheduler.state_dict()							)

           # Load everything back in and make sure all states work
           accelerator.load_state(os.path.join(A_			,					'''checkpoints'''			,					'''checkpoint_0'''							)							)
           self.assertEqual(A_			,					scheduler.state_dict()							)
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         with tempfile.TemporaryDirectory() as tmpdir:
           set_seed(42							)
           _lowercase:      List[Any]   			=  DummyModel()
           _lowercase:      List[Any]   			=  ProjectConfiguration(automatic_checkpoint_naming=A_			,					total_limit=2							)
           # Train baseline
           _lowercase:      Union[str, Any]   			=  Accelerator(project_dir=A_			,					project_config=A_							)
           _lowercase:      Optional[Any]   			=  accelerator.prepare(A_							)
           # Save 3 states:
           for _ in range(11							):
             accelerator.save_state()
           self.assertTrue(not os.path.exists(os.path.join(A_			,					'''checkpoints'''			,					'''checkpoint_0'''							)							)							)
           self.assertTrue(os.path.exists(os.path.join(A_			,					'''checkpoints'''			,					'''checkpoint_9'''							)							)							)
           self.assertTrue(os.path.exists(os.path.join(A_			,					'''checkpoints'''			,					'''checkpoint_10'''							)							)							)



       @require_cuda
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         _lowercase:      List[Any]   			=  ['''torchrun''', f'''--nproc_per_node={torch.cuda.device_count()}''', inspect.getfile(self.__class__							)]
         execute_subprocess_async(A_			,					env=os.environ.copy()							)


if __name__ == "__main__":
  A__ :       Dict				= '/tmp/accelerate/state_checkpointing'
  A__ :       Union[str, Any]				= DummyModel()
  A__ :       str				= torch.optim.Adam(params=model.parameters(), lr=1e-3)
  A__ :       Union[str, Any]				= torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
  A__				,			A__ :       Optional[int]				= dummy_dataloaders()
  A__ :       Union[str, Any]				= ProjectConfiguration(automatic_checkpoint_naming=True)
  # Train baseline
  A__ :       List[Any]				= Accelerator(project_dir=savedir, project_config=project_config, mixed_precision='no')
  if accelerator.process_index == 0:
    if os.path.exists(savedir):
      shutil.rmtree(savedir)
    os.makedirs(savedir)
  A__				,			A__				,			A__				,			A__				,			A__ :       Union[str, Any]				= accelerator.prepare(
      model, optimizer, train_dataloader, valid_dataloader, scheduler
  )
  A__				,			A__ :       str				= accelerator.prepare(model, optimizer)
  train(3, model, train_dataloader, optimizer, accelerator, scheduler)
  # Check that the intial optimizer is loaded on the GPU
  for group in optimizer.param_groups:
    A__ :       Tuple				= group['params'][0].device
    break
  assert param_device.type == accelerator.device.type
  A__ :       Tuple				= model.cpu()
  accelerator.wait_for_everyone()
  accelerator.save_state()
  accelerator.wait_for_everyone()

  # Check CPU state
  accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='cpu')
  for group in optimizer.param_groups:
    A__ :       int				= group['params'][0].device
    break
  assert (
      param_device.type == torch.device('cpu').type
  ), f"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

  # Check device state
  model.to(accelerator.device)
  accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='on_device')
  for group in optimizer.param_groups:
    A__ :       List[Any]				= group['params'][0].device
    break
  assert (
      param_device.type == accelerator.device.type
  ), f"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

  # Check error
  with pytest.raises(TypeError, match='Unsupported optimizer map location passed'):
    accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='invalid')
  accelerator.wait_for_everyone()
  if accelerator.process_index == 0:
    shutil.rmtree(savedir)
  accelerator.wait_for_everyone()


"""simple docstring"""




import functools
import logging
import os
import sys
import threading
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

import huggingface_hub.utils as hf_hub_utils
from tqdm import auto as tqdm_lib


A__ :       int				= threading.Lock()
A__ :       Optional[logging.Handler]				= None

A__ :       str				= {
    'debug': logging.DEBUG,
    'info': logging.INFO,
    'warning': logging.WARNING,
    'error': logging.ERROR,
    'critical': logging.CRITICAL,
}

A__ :       str				= logging.WARNING

A__ :       Union[str, Any]				= True




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _lowercase:      List[Any]   			=  os.getenv('''TRANSFORMERS_VERBOSITY'''			,					_UpperCamelCase    )
  if env_level_str:
    if env_level_str in log_levels:
      return log_levels[env_level_str]
    else:
      logging.getLogger().warning(
          f'''Unknown option TRANSFORMERS_VERBOSITY={env_level_str}, '''
          f'''has to be one of: { ", ".join(log_levels.keys()    ) }'''    )
  return _default_log_level




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  return __name__.split('''.'''    )[0]




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  return logging.getLogger(_get_library_name()    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  global _default_handler

  with _lock:
    if _default_handler:
      # This library has already configured the library root logger.
      return
    _lowercase:      int   			=  logging.StreamHandler()  # Set sys.stderr as stream.
    _lowercase:      Dict   			=  sys.stderr.flush

    # Apply our default configuration to the library root logger.
    _lowercase:      Dict   			=  _get_library_root_logger()
    library_root_logger.addHandler(_default_handler    )
    library_root_logger.setLevel(_get_default_logging_level()    )
    _lowercase:      Optional[Any]   			=  False




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  global _default_handler

  with _lock:
    if not _default_handler:
      return

    _lowercase:      Tuple   			=  _get_library_root_logger()
    library_root_logger.removeHandler(_default_handler    )
    library_root_logger.setLevel(logging.NOTSET    )
    _lowercase:      Dict   			=  None




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  return log_levels




def 			_lowerCAmelCase		(		_UpperCamelCase = None    ):






  """simple docstring"""



  if name is None:
    _lowercase:      Tuple   			=  _get_library_name()

  _configure_library_root_logger()
  return logging.getLogger(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _configure_library_root_logger()
  return _get_library_root_logger().getEffectiveLevel()




def 			_lowerCAmelCase		(		_UpperCamelCase    ):






  """simple docstring"""



  _configure_library_root_logger()
  _get_library_root_logger().setLevel(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  return set_verbosity(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  return set_verbosity(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  return set_verbosity(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  return set_verbosity(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _configure_library_root_logger()

  assert _default_handler is not None
  _get_library_root_logger().removeHandler(_default_handler    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _configure_library_root_logger()

  assert _default_handler is not None
  _get_library_root_logger().addHandler(_default_handler    )




def 			_lowerCAmelCase		(		_UpperCamelCase    ):






  """simple docstring"""



  _configure_library_root_logger()

  assert handler is not None
  _get_library_root_logger().addHandler(_UpperCamelCase    )




def 			_lowerCAmelCase		(		_UpperCamelCase    ):






  """simple docstring"""



  _configure_library_root_logger()

  assert handler is not None and handler not in _get_library_root_logger().handlers
  _get_library_root_logger().removeHandler(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _configure_library_root_logger()
  _lowercase:      str   			=  False




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _configure_library_root_logger()
  _lowercase:      List[str]   			=  True




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _lowercase:      Any   			=  _get_library_root_logger().handlers

  for handler in handlers:
    _lowercase:      int   			=  logging.Formatter('''[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s'''    )
    handler.setFormatter(_UpperCamelCase    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _lowercase:      List[str]   			=  _get_library_root_logger().handlers

  for handler in handlers:
    handler.setFormatter(_UpperCamelCase    )




def 			_lowerCAmelCase		(		self			,					*_UpperCamelCase			,					**_UpperCamelCase    ):






  """simple docstring"""



  _lowercase:      Any   			=  os.getenv('''TRANSFORMERS_NO_ADVISORY_WARNINGS'''			,					_UpperCamelCase    )
  if no_advisory_warnings:
    return
  self.warning(*_UpperCamelCase			,					**_UpperCamelCase    )


A__ :       Optional[int]				= warning_advice




@functools.lru_cache(_UpperCamelCase    )
def 			_lowerCAmelCase		(		self			,					*_UpperCamelCase			,					**_UpperCamelCase    ):






  """simple docstring"""



  self.warning(*_UpperCamelCase			,					**_UpperCamelCase    )


A__ :       List[Any]				= warning_once


class       __magic_name__							:
       def __init__(  self			,					*A_			,					**A_							)  ->						Any:  # pylint: disable=unused-argument





         """simple docstring"""




         _lowercase:      Tuple   			=  args[0] if args else None
       def __iter__(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         return iter(self._iterator							)
       def __getattr__(  self			,					A_							)  ->						List[Any]:





         """simple docstring"""




         def empty_fn(*A_			,					**A_							):  # pylint: disable=unused-argument
           return

         return empty_fn
       def __enter__(  self							)  ->						Tuple:





         """simple docstring"""




         return self



       def __exit__(  self			,					A_			,					A_			,					A_							)  ->						Optional[Any]:





         """simple docstring"""




         return


class       __magic_name__							:
       def __call__(  self			,					*A_			,					**A_							)  ->						Dict:





         """simple docstring"""




         if _tqdm_active:
           return tqdm_lib.tqdm(*A_			,					**A_							)
         else:
           return EmptyTqdm(*A_			,					**A_							)
       def 				lowercase_				(  self			,					*A_			,					**A_							)  ->						List[str]:





         """simple docstring"""




         _lowercase:      Optional[Any]   			=  None
         if _tqdm_active:
           return tqdm_lib.tqdm.set_lock(*A_			,					**A_							)



       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         if _tqdm_active:
           return tqdm_lib.tqdm.get_lock()

A__ :       str				= _tqdm_cls()




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  global _tqdm_active
  return bool(_tqdm_active    )




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  global _tqdm_active
  _lowercase:      str   			=  True
  hf_hub_utils.enable_progress_bars()







def 			_lowerCAmelCase		(		):






  """simple docstring"""



  global _tqdm_active
  _lowercase:      Union[str, Any]   			=  False
  hf_hub_utils.disable_progress_bars()
