




"""simple docstring"""






import gc
import threading
import time

import psutil
import torch






class lowerCAmelCase__	:






			'''simple docstring'''



			def __init__(  self  ):
							_lowerCamelCase     :			Tuple    =    psutil.Process()
							_lowerCamelCase     :			Optional[int]    =    False
			def 					A_    (  self  ):
							_lowerCamelCase     :			Tuple    =    -1

							while True:
											_lowerCamelCase     :			List[Any]    =    max(self.process.memory_info().rss  ,      self.cpu_memory_peak  )

											# can't sleep or will not catch the peak right (this comment is here on purpose)
											if not self.peak_monitoring:
															break
			def 					A_    (  self  ):
							_lowerCamelCase     :			List[str]    =    True
							_lowerCamelCase     :			Union[str, Any]    =    threading.Thread(target=self.peak_monitor  )
							_lowerCamelCase     :			Optional[Any]    =    True
							self.thread.start()






			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    False
							self.thread.join()
							return self.cpu_memory_peak

lowercase__     						=						PeakCPUMemory()







def      _snake_case			(				):
				# Time
				_lowerCamelCase     :			Dict    =    {'time': time.time()}

				gc.collect()
				torch.cuda.empty_cache()

				# CPU mem
				_lowerCamelCase     :			str    =    psutil.Process().memory_info().rss
				cpu_peak_tracker.start()

				# GPU mem
				for i in range(torch.cuda.device_count()		):
								_lowerCamelCase     :			List[Any]    =    torch.cuda.memory_allocated(lowercase__		)
				torch.cuda.reset_peak_memory_stats()

				return measures







def      _snake_case			(				lowercase__		):
				# Time
				_lowerCamelCase     :			Any    =    {'time': time.time() - start_measures['time']}

				gc.collect()
				torch.cuda.empty_cache()

				# CPU mem
				_lowerCamelCase     :			Dict    =    (psutil.Process().memory_info().rss - start_measures['cpu']) / 2**20
				_lowerCamelCase     :			str    =    (cpu_peak_tracker.stop() - start_measures['cpu']) / 2**20

				# GPU mem
				for i in range(torch.cuda.device_count()		):
								_lowerCamelCase     :			List[Any]    =    (torch.cuda.memory_allocated(lowercase__		) - start_measures[str(lowercase__		)]) / 2**20
								_lowerCamelCase     :			Union[str, Any]    =    (torch.cuda.max_memory_allocated(lowercase__		) - start_measures[str(lowercase__		)]) / 2**20

				return measures







def      _snake_case			(				lowercase__							,    lowercase__		):
				print(f'''{description}:'''		)
				print(f'''- Time: {measures['time']:.2f}s'''		)
				for i in range(torch.cuda.device_count()		):
								print(f'''- GPU {i} allocated: {measures[str(lowercase__		)]:.2f}MiB'''		)
								_lowerCamelCase     :			Any    =    measures[f'''{i}-peak''']
								print(f'''- GPU {i} peak: {peak:.2f}MiB'''		)
				print(f'''- CPU RAM allocated: {measures['cpu']:.2f}MiB'''		)
				print(f'''- CPU RAM peak: {measures['cpu-peak']:.2f}MiB'''		)





"""simple docstring"""






import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

lowercase__     						=						logging.getLogger()







def      _snake_case			(				lowercase__		):
				_lowerCamelCase     :			List[Any]    =    {}
				_lowerCamelCase     :			List[Any]    =    os.path.join(lowercase__							,    'all_results.json'		)
				if os.path.exists(lowercase__		):
								with open(lowercase__							,    'r'		) as f:
												_lowerCamelCase     :			List[Any]    =    json.load(lowercase__		)
				else:
								raise ValueError(f'''can\'t find {path}'''		)
				return results


lowercase__     						=						logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)






@require_torch_tpu
class lowerCAmelCase__	(     lowercase						):






			'''simple docstring'''



			def 					A_    (  self  ):
							import xla_spawn

							_lowerCamelCase     :			List[Any]    =    self.get_auto_remove_tmp_dir()
							_lowerCamelCase     :			List[Any]    =    F'''
            ./examples/pytorch/text-classification/run_glue.py
            --num_cores=8
            ./examples/pytorch/text-classification/run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --do_train
            --do_eval
            --debug tpu_metrics_debug
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --max_steps=10
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            '''.split()

							with patch.object(lowercase  ,      'argv'  ,      lowercase  ):
											_lowerCamelCase     :			Dict    =    time()
											xla_spawn.main()
											_lowerCamelCase     :			Any    =    time()

											_lowerCamelCase     :			Optional[int]    =    get_results(lowercase  )
											self.assertGreaterEqual(result['eval_accuracy']  ,      0.75  )

											# Assert that the script takes less than 500 seconds to make sure it doesn't hang.
											self.assertLess(end - start  ,      500  )






			def 					A_    (  self  ):
							import xla_spawn

							_lowerCamelCase     :			Tuple    =    '\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            '.split()
							with patch.object(lowercase  ,      'argv'  ,      lowercase  ):
											xla_spawn.main()