




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)

lowercase_	   =     [
    '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 	A_ (	unittest.TestCase   ):




      '''simple docstring'''


      def 					_snake_case     (    self:	Any						,				a:	str						,				a:	bool						,				a:	str = None						,				a:	list = None    ):
       __lowerCamelCase		:							Optional[Any]		 =						None
       __lowerCamelCase		:							Tuple		 =						os.path.abspath(os.path.join('examples'						,				'by_feature'    )    )
       __lowerCamelCase		:							Dict		 =						os.path.abspath('examples'    )
       for item in os.listdir(a    ):
        if item not in EXCLUDE_EXAMPLES:
         __lowerCamelCase		:							int		 =						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()'						,				):
           __lowerCamelCase		:							str		 =						compare_against_test(
               os.path.join(a						,				a    )						,				a						,				a						,				a    )
           __lowerCamelCase		:							Tuple		 =						'\n'.join(a    )
           if special_strings is not None:
            for string in special_strings:
             __lowerCamelCase		:							Dict		 =						diff.replace(a						,				''    )
           self.assertEqual(a						,				''    )




      def 					_snake_case     (    self:	int    ):
       self.one_complete_example('complete_nlp_example.py'						,				a    )
       self.one_complete_example('complete_nlp_example.py'						,				a    )




      def 					_snake_case     (    self:	Optional[Any]    ):
       __lowerCamelCase		:							List[Any]		 =						os.path.abspath(os.path.join('examples'						,				'cv_example.py'    )    )
       __lowerCamelCase		:							Optional[int]		 =						[
           ' ' * 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 	A_ (	__UpperCamelCase   ):




      '''simple docstring'''


      __snake_case 	=       False




      @classmethod
      def 					_snake_case     (    cls:	List[str]    ):
       super().setUpClass()
       __lowerCamelCase		:							str		 =						tempfile.mkdtemp()
       __lowerCamelCase		:							Any		 =						os.path.join(cls._tmpdir						,				'default_config.yml'    )

       write_basic_config(save_location=cls.configPath    )
       __lowerCamelCase		:							int		 =						['accelerate', 'launch', '--config_file', cls.configPath]




      @classmethod
      def 					_snake_case     (    cls:	Optional[int]    ):
       super().tearDownClass()
       shutil.rmtree(cls._tmpdir    )




      def 					_snake_case     (    self:	Union[str, Any]    ):
       __lowerCamelCase		:							Optional[int]		 =						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 					_snake_case     (    self:	Tuple    ):
       __lowerCamelCase		:							Optional[int]		 =						F'\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps 1\n        --output_dir {self.tmpdir}\n        '.split()
       __lowerCamelCase		:							List[str]		 =						run_command(self._launch_args + testargs    )
       self.assertTrue(os.path.exists(os.path.join(self.tmpdir						,				'step_2'    )    )    )




      def 					_snake_case     (    self:	Dict    ):
       __lowerCamelCase		:							Optional[int]		 =						F'\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir						,				"epoch_0"    )}\n        '.split()
       __lowerCamelCase		:							Optional[Any]		 =						run_command(self._launch_args + testargs						,				return_stdout=a    )
       self.assertNotIn('epoch 0:'						,				a    )
       self.assertIn('epoch 1:'						,				a    )




      def 					_snake_case     (    self:	List[str]    ):
       __lowerCamelCase		:							Tuple		 =						F'\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir						,				"step_2"    )}\n        '.split()
       __lowerCamelCase		:							Optional[Any]		 =						run_command(self._launch_args + testargs						,				return_stdout=a    )
       if torch.cuda.is_available():
        __lowerCamelCase		:							Union[str, Any]		 =						torch.cuda.device_count()
       else:
        __lowerCamelCase		:							str		 =						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 					_snake_case     (    self:	Any    ):
       __lowerCamelCase		:							Optional[int]		 =						'\n        examples/by_feature/cross_validation.py\n        --num_folds 2\n        '.split()
       with mock.patch.dict(os.environ						,				{'TESTING_MOCKED_DATALOADERS': '0'}    ):
        __lowerCamelCase		:							Optional[int]		 =						run_command(self._launch_args + testargs						,				return_stdout=a    )
        __lowerCamelCase		:							Union[str, Any]		 =						re.findall('({.+})'						,				a    )
        __lowerCamelCase		:							Dict		 =						[r for r in results if 'accuracy' in r][-1]
        __lowerCamelCase		:							str		 =						ast.literal_eval(a    )
        self.assertGreaterEqual(results['accuracy']						,				0.7_5    )




      def 					_snake_case     (    self:	Optional[Any]    ):
       __lowerCamelCase		:							List[Any]		 =						['examples/by_feature/multi_process_metrics.py']
       run_command(self._launch_args + testargs    )




      @require_trackers
      @mock.patch.dict(os.environ						,				{'WANDB_MODE': 'offline'}    )
      def 					_snake_case     (    self:	Tuple    ):
       with tempfile.TemporaryDirectory() as tmpdir:
        __lowerCamelCase		:							Union[str, Any]		 =						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 					_snake_case     (    self:	Any    ):
       __lowerCamelCase		:							Optional[Any]		 =						['examples/by_feature/gradient_accumulation.py']
       run_command(self._launch_args + testargs    )




      def 					_snake_case     (    self:	Union[str, Any]    ):
       __lowerCamelCase		:							str		 =						['examples/by_feature/local_sgd.py']
       run_command(self._launch_args + testargs    )





import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


lowercase_	   =     5_0_0_0_0
lowercase_	   =     5_0_0_0

lowercase_   ,lowercase_	   =     os.path.split(__file__)
lowercase_	   =     os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 for i in range(SCREAMING_SNAKE_CASE__						):
  __lowerCamelCase		:							Tuple		 =						dataset[i]
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 for i in range(0      ,	len(SCREAMING_SNAKE_CASE__						)      ,	SCREAMING_SNAKE_CASE__						):
  __lowerCamelCase		:							Optional[Any]		 =						dataset[i : i + batch_size]
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 with dataset.formatted_as(type=SCREAMING_SNAKE_CASE__						):
  for i in range(SCREAMING_SNAKE_CASE__						):
   __lowerCamelCase		:							str		 =						dataset[i]
@get_duration
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 with dataset.formatted_as(type=SCREAMING_SNAKE_CASE__						):
  for i in range(0      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
   __lowerCamelCase		:							int		 =						dataset[i : i + batch_size]
def 				UpperCamelCase__					(  ):
 __lowerCamelCase		:							Union[str, Any]		 =						{'num examples': SPEED_TEST_N_EXAMPLES}
 __lowerCamelCase		:							Optional[Any]		 =						[
     (read, {'length': SMALL_TEST}),
     (read, {'length': SPEED_TEST_N_EXAMPLES}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 10}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 100}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 1_000}),
     (read_formatted, {'type': 'numpy', 'length': SMALL_TEST}),
     (read_formatted, {'type': 'pandas', 'length': SMALL_TEST}),
     (read_formatted, {'type': 'torch', 'length': SMALL_TEST}),
     (read_formatted, {'type': 'tensorflow', 'length': SMALL_TEST}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 10}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 1_000}),
 ]

 __lowerCamelCase		:							Any		 =						[
     (read, {'length': SMALL_TEST}),
     (read, {'length': SPEED_TEST_N_EXAMPLES}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 10}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 100}),
     (read_batch, {'length': SPEED_TEST_N_EXAMPLES, 'batch_size': 1_000}),
     (read_formatted, {'type': 'numpy', 'length': SMALL_TEST}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 10}),
     (read_formatted_batch, {'type': 'numpy', 'length': SMALL_TEST, 'batch_size': 1_000}),
 ]
 with tempfile.TemporaryDirectory() as tmp_dir:
  print('generating dataset'						)
  __lowerCamelCase		:							Optional[int]		 =						datasets.Features(
      {'list': datasets.Sequence(datasets.Value('float32'						)						), 'numbers': datasets.Value('float32'						)}						)
  __lowerCamelCase		:							str		 =						generate_example_dataset(
      os.path.join(SCREAMING_SNAKE_CASE__      ,	'dataset.arrow'						)      ,	SCREAMING_SNAKE_CASE__      ,	num_examples=SCREAMING_SNAKE_CASE__      ,	seq_shapes={'list': (100,)}      ,	)
  print('first set of iterations'						)
  for func, kwargs in functions:
   print(func.__name__      ,	str(SCREAMING_SNAKE_CASE__						)						)
   __lowerCamelCase		:							Optional[int]		 =						func(SCREAMING_SNAKE_CASE__      ,	**SCREAMING_SNAKE_CASE__						)

  print('shuffling dataset'						)
  __lowerCamelCase		:							str		 =						dataset.shuffle()
  print('Second set of iterations (after shuffling'						)
  for func, kwargs in functions_shuffled:
   print('shuffled '      ,	func.__name__      ,	str(SCREAMING_SNAKE_CASE__						)						)
   __lowerCamelCase		:							int		 =						func(
       SCREAMING_SNAKE_CASE__      ,	**SCREAMING_SNAKE_CASE__						)

 with open(SCREAMING_SNAKE_CASE__      ,	'wb'						) as f:
  f.write(json.dumps(SCREAMING_SNAKE_CASE__						).encode('utf-8'						)						)


if __name__ == "__main__":  # useful to run the profiler
       benchmark_iterating()

