



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_    :  str						  =		logging.getLogger(__name__)

def      __lowercase(						UpperCAmelCase__=2 , UpperCAmelCase__=3 , UpperCAmelCase__=16 , UpperCAmelCase__ = 10 , UpperCAmelCase__ = 2			):


      """simple docstring"""



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

      lowerCamelCase             =		get_dataset(UpperCAmelCase__			)
      lowerCamelCase             =		get_dataset(UpperCAmelCase__			)
      lowerCamelCase             =		DataLoader(UpperCAmelCase__ , shuffle=UpperCAmelCase__ , batch_size=UpperCAmelCase__ , num_workers=4			)
      lowerCamelCase             =		DataLoader(UpperCAmelCase__ , shuffle=UpperCAmelCase__ , batch_size=UpperCAmelCase__ , num_workers=4			)
      return (train_dataloader, valid_dataloader)

def      __lowercase(						UpperCAmelCase__ , UpperCAmelCase__ , UpperCAmelCase__ , UpperCAmelCase__ , UpperCAmelCase__ , UpperCAmelCase__=None			):


      """simple docstring"""



      lowerCamelCase             =		[]
      for epoch in range(UpperCAmelCase__			):
            # Train quickly
            model.train()
            for batch in dataloader:
                  lowerCamelCase      ,				lowerCamelCase             =		batch
                  lowerCamelCase             =		model(UpperCAmelCase__			)
                  lowerCamelCase             =		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  lowerCamelCase__	(    nn.Module):
     """simple docstring"""



     def __init__(self     ):






           '''simple docstring'''



           super().__init__()
           lowerCamelCase             =		nn.Parameter(torch.randn(1     )     )
           lowerCamelCase             =		nn.Parameter(torch.randn(1     )     )







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           return x * self.a + self.b





class  lowerCamelCase__	(    unittest.TestCase):
     """simple docstring"""



     def     _a   (self     ):






           '''simple docstring'''



           with tempfile.TemporaryDirectory() as tmpdir:
                 set_seed(42     )
                 lowerCamelCase             =		DummyModel()
                 lowerCamelCase             =		torch.optim.Adam(params=model.parameters()					,      lr=1E-3     )
                 lowerCamelCase      ,				lowerCamelCase             =		dummy_dataloaders()
                 lowerCamelCase             =		ProjectConfiguration(total_limit=1					,      project_dir=__a					,      automatic_checkpoint_naming=__a     )
                 # Train baseline
                 lowerCamelCase             =		Accelerator(project_config=__a     )
                 lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase             =		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     _a   (self     ):






           '''simple docstring'''



           with tempfile.TemporaryDirectory() as tmpdir:
                 set_seed(42     )
                 lowerCamelCase             =		DummyModel()
                 lowerCamelCase             =		torch.optim.Adam(params=model.parameters()					,      lr=1E-3     )
                 lowerCamelCase      ,				lowerCamelCase             =		dummy_dataloaders()
                 # Train baseline
                 lowerCamelCase             =		Accelerator()
                 lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase             =		accelerator.prepare(
                     __a					,      __a					,      __a					,      __a     )
                 # Save initial
                 lowerCamelCase             =		os.path.join(__a					,      "initial"     )
                 accelerator.save_state(__a     )
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()
                 lowerCamelCase             =		train(3					,      __a					,      __a					,      __a					,      __a     )
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()

                 # Train partially
                 set_seed(42     )
                 lowerCamelCase             =		DummyModel()
                 lowerCamelCase             =		torch.optim.Adam(params=model.parameters()					,      lr=1E-3     )
                 lowerCamelCase      ,				lowerCamelCase             =		dummy_dataloaders()
                 lowerCamelCase             =		Accelerator()
                 lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase             =		accelerator.prepare(
                     __a					,      __a					,      __a					,      __a     )
                 accelerator.load_state(__a     )
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )

                 lowerCamelCase             =		train(2					,      __a					,      __a					,      __a					,      __a     )
                 # Save everything
                 lowerCamelCase             =		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     )
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )







     def     _a   (self     ):






           '''simple docstring'''



           with tempfile.TemporaryDirectory() as tmpdir:
                 set_seed(42     )
                 lowerCamelCase             =		DummyModel()
                 lowerCamelCase             =		torch.optim.Adam(params=model.parameters()					,      lr=1E-3     )
                 lowerCamelCase      ,				lowerCamelCase             =		dummy_dataloaders()
                 lowerCamelCase             =		ProjectConfiguration(automatic_checkpoint_naming=__a     )

                 # Train baseline
                 lowerCamelCase             =		Accelerator(project_dir=__a					,      project_config=__a     )
                 lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase             =		accelerator.prepare(
                     __a					,      __a					,      __a					,      __a     )
                 # Save initial
                 accelerator.save_state()
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()
                 lowerCamelCase             =		train(3					,      __a					,      __a					,      __a					,      __a     )
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()

                 # Train partially
                 set_seed(42     )
                 lowerCamelCase             =		DummyModel()
                 lowerCamelCase             =		torch.optim.Adam(params=model.parameters()					,      lr=1E-3     )
                 lowerCamelCase      ,				lowerCamelCase             =		dummy_dataloaders()
                 lowerCamelCase             =		ProjectConfiguration(iteration=1					,      automatic_checkpoint_naming=__a     )
                 lowerCamelCase             =		Accelerator(project_dir=__a					,      project_config=__a     )
                 lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase             =		accelerator.prepare(
                     __a					,      __a					,      __a					,      __a     )
                 accelerator.load_state(os.path.join(__a					,      "checkpoints"					,      "checkpoint_0"     )     )
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )

                 lowerCamelCase             =		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     )
                 ((lowerCamelCase)      ,				(lowerCamelCase))             =		model.a.item(), model.b.item()
                 lowerCamelCase             =		optimizer.state_dict()
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )
                 self.assertEqual(__a					,      __a     )







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		torch.tensor([1, 2, 3]     )
           lowerCamelCase             =		torch.tensor([2, 3, 4]     )
           lowerCamelCase             =		DummyModel()
           lowerCamelCase             =		torch.optim.Adam(net.parameters()     )
           lowerCamelCase             =		Accelerator()
           with self.assertRaises(__a     ) as ve:
                 accelerator.register_for_checkpointing(__a					,      __a					,      __a					,      __a     )
           lowerCamelCase             =		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     _a   (self     ):






           '''simple docstring'''



           with tempfile.TemporaryDirectory() as tmpdir:
                 set_seed(42     )
                 lowerCamelCase             =		DummyModel()
                 lowerCamelCase             =		torch.optim.Adam(params=model.parameters()					,      lr=1E-3     )
                 lowerCamelCase             =		torch.optim.lr_scheduler.StepLR(__a					,      step_size=1					,      gamma=0.99     )
                 lowerCamelCase      ,				lowerCamelCase             =		dummy_dataloaders()
                 lowerCamelCase             =		ProjectConfiguration(automatic_checkpoint_naming=__a     )
                 # Train baseline
                 lowerCamelCase             =		Accelerator(project_dir=__a					,      project_config=__a     )
                 lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase      ,				lowerCamelCase             =		accelerator.prepare(
                     __a					,      __a					,      __a					,      __a					,      __a     )
                 # Save initial
                 accelerator.save_state()
                 lowerCamelCase             =		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     _a   (self     ):






           '''simple docstring'''



           with tempfile.TemporaryDirectory() as tmpdir:
                 set_seed(42     )
                 lowerCamelCase             =		DummyModel()
                 lowerCamelCase             =		ProjectConfiguration(automatic_checkpoint_naming=__a					,      total_limit=2     )
                 # Train baseline
                 lowerCamelCase             =		Accelerator(project_dir=__a					,      project_config=__a     )
                 lowerCamelCase             =		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     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		["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_    :  Tuple						  =		torch.optim.Adam(params=model.parameters(), lr=1E-3)
     a_    :  int						  =		torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
     a_  ,			a_    :  Dict						  =		dummy_dataloaders()
     a_    :  Optional[Any]						  =		ProjectConfiguration(automatic_checkpoint_naming=True)
     # Train baseline
     a_    :  Dict						  =		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_    :  List[str]						  =		accelerator.prepare(
         model, optimizer, train_dataloader, valid_dataloader, scheduler
     )
     a_  ,			a_    :  Optional[int]						  =		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_    :  Any						  =		group['params'][0].device
          break
     assert param_device.type == accelerator.device.type
     a_    :  List[Any]						  =		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_    :  Dict						  =		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[str]						  =		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()




from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_tf_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_tf_available():
     import tensorflow as tf


a_    :  Any						  =		logging.get_logger(__name__)





@dataclass
class  lowerCamelCase__	(    UpperCAmelCase_):
     """simple docstring"""



     _A      =      [
         'no_inference',
         'no_cuda',
         'no_tpu',
         'no_speed',
         'no_memory',
         'no_env_print',
         'no_multi_process',
     ]







     def __init__(self					,      **__a     ):






           '''simple docstring'''



           for deprecated_arg in self.deprecated_args:
                 if deprecated_arg in kwargs:
                       lowerCamelCase             =		deprecated_arg[3:]
                       lowerCamelCase             =		not kwargs.pop(__a     )
                       logger.warning(
                           F"""{deprecated_arg} is depreciated. Please use --no-{positive_arg} or"""
                           F""" {positive_arg}={kwargs[positive_arg]}"""     )
           lowerCamelCase             =		kwargs.pop("tpu_name"					,      self.tpu_name     )
           lowerCamelCase             =		kwargs.pop("device_idx"					,      self.device_idx     )
           lowerCamelCase             =		kwargs.pop("eager_mode"					,      self.eager_mode     )
           lowerCamelCase             =		kwargs.pop("use_xla"					,      self.use_xla     )
           super().__init__(**__a     )

     _A      =      field(
         default=UpperCAmelCase_							,     metadata={'help': 'Name of TPU'}							,     )
     _A      =      field(
         default=0							,     metadata={'help': 'CPU / GPU device index. Defaults to 0.'}							,     )
     _A      =      field(default=UpperCAmelCase_							,     metadata={'help': 'Benchmark models in eager model.'})
     _A      =      field(
         default=UpperCAmelCase_							,     metadata={
             'help': 'Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`.'
         }							,     )







     @cached_property
     def     _a   (self     ):






           '''simple docstring'''



           requires_backends(self					,      ["tf"]     )
           lowerCamelCase             =		None
           if self.tpu:
                 try:
                       if self.tpu_name:
                             lowerCamelCase             =		tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name     )
                       else:
                             lowerCamelCase             =		tf.distribute.cluster_resolver.TPUClusterResolver()
                 except ValueError:
                       lowerCamelCase             =		None
           return tpu







     @cached_property
     def     _a   (self     ):






           '''simple docstring'''



           requires_backends(self					,      ["tf"]     )
           if self.is_tpu:
                 tf.config.experimental_connect_to_cluster(self._setup_tpu     )
                 tf.tpu.experimental.initialize_tpu_system(self._setup_tpu     )

                 lowerCamelCase             =		tf.distribute.TPUStrategy(self._setup_tpu     )
           else:
                 # currently no multi gpu is allowed
                 if self.is_gpu:
                       # TODO: Currently only single GPU is supported
                       tf.config.set_visible_devices(self.gpu_list[self.device_idx]					,      "GPU"     )
                       lowerCamelCase             =		tf.distribute.OneDeviceStrategy(device=F"""/gpu:{self.device_idx}"""     )
                 else:
                       tf.config.set_visible_devices([]					,      "GPU"     )  # disable GPU
                       lowerCamelCase             =		tf.distribute.OneDeviceStrategy(device=F"""/cpu:{self.device_idx}"""     )

           return strategy







     @property
     def     _a   (self     ):






           '''simple docstring'''



           requires_backends(self					,      ["tf"]     )
           return self._setup_tpu is not None







     @property
     def     _a   (self     ):






           '''simple docstring'''



           requires_backends(self					,      ["tf"]     )
           return self._setup_strategy







     @property
     def     _a   (self     ):






           '''simple docstring'''



           requires_backends(self					,      ["tf"]     )
           return tf.config.list_physical_devices("GPU"     )







     @property
     def     _a   (self     ):






           '''simple docstring'''



           requires_backends(self					,      ["tf"]     )
           if self.cuda:
                 return len(self.gpu_list     )
           return 0







     @property
     def     _a   (self     ):






           '''simple docstring'''



           return self.n_gpu > 0