

import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

UpperCamelCase__  :      Union[str, Any]			=    [
    'cross_validation.py',
    'gradient_accumulation.py',
    'local_sgd.py',
    'multi_process_metrics.py',
    'memory.py',
    'automatic_gradient_accumulation.py',
    'fsdp_with_peak_mem_tracking.py',
    'deepspeed_with_config_support.py',
    'megatron_lm_gpt_pretraining.py',
]



class       __snake_case  (    unittest.TestCase       ):

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A = None						,		_A = None):
      SCREAMING_SNAKE_CASE_									=						None
      SCREAMING_SNAKE_CASE_									=						os.path.abspath(os.path.join('examples'						,		'by_feature'))
      SCREAMING_SNAKE_CASE_									=						os.path.abspath('examples')
      for item in os.listdir(_A):
          if item not in EXCLUDE_EXAMPLES:
              SCREAMING_SNAKE_CASE_									=						os.path.join(_A						,		_A)
              if os.path.isfile(_A) and ".py" in item_path:
                  with self.subTest(
                      tested_script=_A						,		feature_script=_A						,		tested_section='main()' if parser_only else 'training_function()'						,		):
                      SCREAMING_SNAKE_CASE_									=						compare_against_test(
                          os.path.join(_A						,		_A)						,		_A						,		_A						,		_A)
                      SCREAMING_SNAKE_CASE_									=						"\n".join(_A)
                      if special_strings is not None:
                          for string in special_strings:
                              SCREAMING_SNAKE_CASE_									=						diff.replace(_A						,		'')
                      self.assertEqual(_A						,		'')

  def 						lowerCAmelCase__     (					self):
      self.one_complete_example('complete_nlp_example.py'						,		_A)
      self.one_complete_example('complete_nlp_example.py'						,		_A)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						os.path.abspath(os.path.join('examples'						,		'cv_example.py'))
      SCREAMING_SNAKE_CASE_									=						[
          " " * 16 + "{\n\n",
          " " * 20 + "\"accuracy\": eval_metric[\"accuracy\"],\n\n",
          " " * 20 + "\"f1\": eval_metric[\"f1\"],\n\n",
          " " * 20 + "\"train_loss\": total_loss.item() / len(train_dataloader),\n\n",
          " " * 20 + "\"epoch\": epoch,\n\n",
          " " * 16 + "},\n\n",
          " " * 16 + "step=epoch,\n",
          " " * 12,
          " " * 8 + "for step, batch in enumerate(active_dataloader):\n",
      ]
      self.one_complete_example('complete_cv_example.py'						,		_A						,		_A						,		_A)
      self.one_complete_example('complete_cv_example.py'						,		_A						,		_A						,		_A)






@mock.patch.dict(os.environ , {'TESTING_MOCKED_DATALOADERS': '1'}       )
class       __snake_case  (    __lowercase       ):
  __lowerCAmelCase     :	Tuple						 =  False

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      super().setUpClass()
      SCREAMING_SNAKE_CASE_									=						tempfile.mkdtemp()
      SCREAMING_SNAKE_CASE_									=						os.path.join(cls._tmpdir						,		'default_config.yml')

      write_basic_config(save_location=cls.configPath)
      SCREAMING_SNAKE_CASE_									=						["accelerate", "launch", "--config_file", cls.configPath]

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      super().tearDownClass()
      shutil.rmtree(cls._tmpdir)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						f"""\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps epoch\n        --output_dir {self.tmpdir}\n        """.split()
      run_command(self._launch_args + testargs)
      self.assertTrue(os.path.exists(os.path.join(self.tmpdir						,		'epoch_0')))

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						f"""\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps 1\n        --output_dir {self.tmpdir}\n        """.split()
      SCREAMING_SNAKE_CASE_									=						run_command(self._launch_args + testargs)
      self.assertTrue(os.path.exists(os.path.join(self.tmpdir						,		'step_2')))

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						f"""\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir						,		'epoch_0')}\n        """.split()
      SCREAMING_SNAKE_CASE_									=						run_command(self._launch_args + testargs						,		return_stdout=_A)
      self.assertNotIn('epoch 0:'						,		_A)
      self.assertIn('epoch 1:'						,		_A)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						f"""\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir						,		'step_2')}\n        """.split()
      SCREAMING_SNAKE_CASE_									=						run_command(self._launch_args + testargs						,		return_stdout=_A)
      if torch.cuda.is_available():
          SCREAMING_SNAKE_CASE_									=						torch.cuda.device_count()
      else:
          SCREAMING_SNAKE_CASE_									=						1
      if num_processes > 1:
          self.assertNotIn('epoch 0:'						,		_A)
          self.assertIn('epoch 1:'						,		_A)
      else:
          self.assertIn('epoch 0:'						,		_A)
          self.assertIn('epoch 1:'						,		_A)

  @slow
  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						"\n        examples/by_feature/cross_validation.py\n        --num_folds 2\n        ".split()
      with mock.patch.dict(os.environ						,		{'TESTING_MOCKED_DATALOADERS': '0'}):
          SCREAMING_SNAKE_CASE_									=						run_command(self._launch_args + testargs						,		return_stdout=_A)
          SCREAMING_SNAKE_CASE_									=						re.findall('({.+})'						,		_A)
          SCREAMING_SNAKE_CASE_									=						[r for r in results if "accuracy" in r][-1]
          SCREAMING_SNAKE_CASE_									=						ast.literal_eval(_A)
          self.assertGreaterEqual(results['accuracy']						,		0.7_5)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						["examples/by_feature/multi_process_metrics.py"]
      run_command(self._launch_args + testargs)

  @require_trackers
  @mock.patch.dict(os.environ						,		{'WANDB_MODE': 'offline'})
  def 						lowerCAmelCase__     (					self):
      with tempfile.TemporaryDirectory() as tmpdir:
          SCREAMING_SNAKE_CASE_									=						f"""\n            examples/by_feature/tracking.py\n            --with_tracking\n            --project_dir {tmpdir}\n            """.split()
          run_command(self._launch_args + testargs)
          self.assertTrue(os.path.exists(os.path.join(_A						,		'tracking')))

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						["examples/by_feature/gradient_accumulation.py"]
      run_command(self._launch_args + testargs)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						["examples/by_feature/local_sgd.py"]
      run_command(self._launch_args + testargs)





import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command



class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Dict						 =  inspect.getfile(accelerate.test_utils       )
  __lowerCAmelCase     :	Optional[Any]						 =  os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_cli.py']       )

  __lowerCAmelCase     :	Tuple						 =  ['accelerate', 'launch']
  __lowerCAmelCase     :	Union[str, Any]						 =  Path.home() / '.cache/huggingface/accelerate'
  __lowerCAmelCase     :	List[str]						 =  'default_config.yaml'
  __lowerCAmelCase     :	List[Any]						 =  config_folder / config_file
  __lowerCAmelCase     :	str						 =  config_folder / '_default_config.yaml'

  __lowerCAmelCase     :	Optional[int]						 =  Path('tests/test_configs'       )

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.config_path.is_file():
          cls.config_path.rename(cls.changed_path)

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.changed_path.is_file():
          cls.changed_path.rename(cls.config_path)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.base_cmd
      if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
          cmd += ["--multi_gpu"]
      execute_subprocess_async(cmd + [self.test_file_path]						,		env=os.environ.copy())

  def 						lowerCAmelCase__     (					self):
      for config in sorted(self.test_config_path.glob('**/*.yaml')):
          with self.subTest(config_file=_A):
              execute_subprocess_async(
                  self.base_cmd + ['--config_file', str(_A), self.test_file_path]						,		env=os.environ.copy())


  def 						lowerCAmelCase__     (					self):
      execute_subprocess_async(['accelerate', 'test']						,		env=os.environ.copy())






class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Optional[Any]						 =  'test-tpu'
  __lowerCAmelCase     :	str						 =  'us-central1-a'
  __lowerCAmelCase     :	Union[str, Any]						 =  'ls'
  __lowerCAmelCase     :	Union[str, Any]						 =  ['accelerate', 'tpu-config']
  __lowerCAmelCase     :	Union[str, Any]						 =  'cd /usr/share'
  __lowerCAmelCase     :	List[Any]						 =  'tests/test_samples/test_command_file.sh'
  __lowerCAmelCase     :	Dict						 =  'Running gcloud compute tpus tpu-vm ssh'

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--command', self.command, '--tpu_zone', self.tpu_zone, '--tpu_name', self.tpu_name, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command',
              self.command,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--debug']						,		return_stdout=_A)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--command', self.command, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--command',
              self.command,
              '--command',
              'echo "Hello World"',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo \"Hello World\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--config_file', 'tests/test_configs/latest.yaml', '--command_file', self.command_file, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command_file',
              self.command_file,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--install_accelerate', '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--install_accelerate',
              '--accelerate_version',
              '12.0.0',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


