"""simple docstring"""





import shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
				from transformers.trainer import DEFAULT_CALLBACKS

				from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel
class  _snake_case     (					a__		):






						def __init__(       self      :			Tuple    ):
													__lowerCamelCase					:		Any		   =     []






						def       lowerCamelCase__     (       self      :			Optional[Any]							,					UpperCAmelCase      :			Optional[Any]							,					UpperCAmelCase      :			str							,					UpperCAmelCase      :			Optional[Any]							,					**UpperCAmelCase      :			Tuple    ):
													self.events.append("on_init_end"    )






						def       lowerCamelCase__     (       self      :			List[Any]							,					UpperCAmelCase      :			List[Any]							,					UpperCAmelCase      :			int							,					UpperCAmelCase      :			List[Any]							,					**UpperCAmelCase      :			Union[str, Any]    ):
													self.events.append("on_train_begin"    )






						def       lowerCamelCase__     (       self      :			List[str]							,					UpperCAmelCase      :			str							,					UpperCAmelCase      :			List[str]							,					UpperCAmelCase      :			List[str]							,					**UpperCAmelCase      :			Optional[Any]    ):
													self.events.append("on_train_end"    )






						def       lowerCamelCase__     (       self      :			List[str]							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			Tuple							,					UpperCAmelCase      :			str							,					**UpperCAmelCase      :			int    ):
													self.events.append("on_epoch_begin"    )






						def       lowerCamelCase__     (       self      :			Dict							,					UpperCAmelCase      :			List[Any]							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			Dict							,					**UpperCAmelCase      :			List[Any]    ):
													self.events.append("on_epoch_end"    )






						def       lowerCamelCase__     (       self      :			Optional[int]							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			List[Any]							,					UpperCAmelCase      :			Tuple							,					**UpperCAmelCase      :			Union[str, Any]    ):
													self.events.append("on_step_begin"    )






						def       lowerCamelCase__     (       self      :			str							,					UpperCAmelCase      :			int							,					UpperCAmelCase      :			int							,					UpperCAmelCase      :			Optional[Any]							,					**UpperCAmelCase      :			List[str]    ):
													self.events.append("on_step_end"    )






						def       lowerCamelCase__     (       self      :			Union[str, Any]							,					UpperCAmelCase      :			Optional[Any]							,					UpperCAmelCase      :			Optional[int]							,					UpperCAmelCase      :			Dict							,					**UpperCAmelCase      :			Tuple    ):
													self.events.append("on_evaluate"    )






						def       lowerCamelCase__     (       self      :			Any							,					UpperCAmelCase      :			List[str]							,					UpperCAmelCase      :			str							,					UpperCAmelCase      :			Optional[Any]							,					**UpperCAmelCase      :			List[str]    ):
													self.events.append("on_predict"    )






						def       lowerCamelCase__     (       self      :			Dict							,					UpperCAmelCase      :			Dict							,					UpperCAmelCase      :			Union[str, Any]							,					UpperCAmelCase      :			str							,					**UpperCAmelCase      :			Tuple    ):
													self.events.append("on_save"    )






						def       lowerCamelCase__     (       self      :			List[str]							,					UpperCAmelCase      :			int							,					UpperCAmelCase      :			Dict							,					UpperCAmelCase      :			int							,					**UpperCAmelCase      :			str    ):
													self.events.append("on_log"    )







						def       lowerCamelCase__     (       self      :			List[Any]							,					UpperCAmelCase      :			Optional[Any]							,					UpperCAmelCase      :			Optional[int]							,					UpperCAmelCase      :			Optional[Any]							,					**UpperCAmelCase      :			Optional[Any]    ):
													self.events.append("on_prediction_step"    )

@require_torch
class  _snake_case     (					unittest.TestCase		):






						def       lowerCamelCase__     (       self      :			str    ):
													__lowerCamelCase					:		List[Any]		   =     tempfile.mkdtemp()






						def       lowerCamelCase__     (       self      :			List[str]    ):
													shutil.rmtree(self.output_dir    )






						def       lowerCamelCase__     (       self      :			Optional[int]							,					UpperCAmelCase      :			List[Any]=0							,					UpperCAmelCase      :			Optional[Any]=0							,					UpperCAmelCase      :			Any=64							,					UpperCAmelCase      :			Union[str, Any]=64							,					UpperCAmelCase      :			str=None							,					UpperCAmelCase      :			Tuple=False							,					**UpperCAmelCase      :			Union[str, Any]    ):
													# disable_tqdm in TrainingArguments has a flaky default since it depends on the level of logging. We make sure
													# its set to False since the tests later on depend on its value.
													__lowerCamelCase					:		Optional[int]		   =     RegressionDataset(length=UpperCAmelCase    )
													__lowerCamelCase					:		int		   =     RegressionDataset(length=UpperCAmelCase    )
													__lowerCamelCase					:		str		   =     RegressionModelConfig(a=UpperCAmelCase							,					b=UpperCAmelCase    )
													__lowerCamelCase					:		List[str]		   =     RegressionPreTrainedModel(UpperCAmelCase    )

													__lowerCamelCase					:		str		   =     TrainingArguments(self.output_dir							,					disable_tqdm=UpperCAmelCase							,					report_to=[]							,					**UpperCAmelCase    )
													return Trainer(
													    UpperCAmelCase							,					UpperCAmelCase							,					train_dataset=UpperCAmelCase							,					eval_dataset=UpperCAmelCase							,					callbacks=UpperCAmelCase							,					)






						def       lowerCamelCase__     (       self      :			Optional[Any]							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			Union[str, Any]    ):
													self.assertEqual(len(UpperCAmelCase    )							,					len(UpperCAmelCase    )    )

													# Order doesn't matter
													__lowerCamelCase					:		str		   =     sorted(UpperCAmelCase							,					key=lambda UpperCAmelCase    : cb.__name__ if isinstance(UpperCAmelCase							,					UpperCAmelCase    ) else cb.__class__.__name__    )
													__lowerCamelCase					:		Union[str, Any]		   =     sorted(UpperCAmelCase							,					key=lambda UpperCAmelCase    : cb.__name__ if isinstance(UpperCAmelCase							,					UpperCAmelCase    ) else cb.__class__.__name__    )

													for cba, cba in zip(UpperCAmelCase							,					UpperCAmelCase    ):
																				if isinstance(UpperCAmelCase							,					UpperCAmelCase    ) and isinstance(UpperCAmelCase							,					UpperCAmelCase    ):
																											self.assertEqual(UpperCAmelCase							,					UpperCAmelCase    )
																				elif isinstance(UpperCAmelCase							,					UpperCAmelCase    ) and not isinstance(UpperCAmelCase							,					UpperCAmelCase    ):
																											self.assertEqual(UpperCAmelCase							,					cba.__class__    )
																				elif not isinstance(UpperCAmelCase							,					UpperCAmelCase    ) and isinstance(UpperCAmelCase							,					UpperCAmelCase    ):
																											self.assertEqual(cba.__class__							,					UpperCAmelCase    )
																				else:
																											self.assertEqual(UpperCAmelCase							,					UpperCAmelCase    )






						def       lowerCamelCase__     (       self      :			str							,					UpperCAmelCase      :			Tuple    ):
													__lowerCamelCase					:		Tuple		   =     ["on_init_end", "on_train_begin"]
													__lowerCamelCase					:		Union[str, Any]		   =     0
													__lowerCamelCase					:		Any		   =     len(trainer.get_eval_dataloader()    )
													__lowerCamelCase					:		List[Any]		   =     ["on_prediction_step"] * len(trainer.get_eval_dataloader()    ) + ["on_log", "on_evaluate"]
													for _ in range(trainer.state.num_train_epochs    ):
																				expected_events.append("on_epoch_begin"    )
																				for _ in range(UpperCAmelCase    ):
																											step += 1
																											expected_events += ["on_step_begin", "on_step_end"]
																											if step % trainer.args.logging_steps == 0:
																																		expected_events.append("on_log"    )
																											if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
																																		expected_events += evaluation_events.copy()
																											if step % trainer.args.save_steps == 0:
																																		expected_events.append("on_save"    )
																				expected_events.append("on_epoch_end"    )
																				if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
																											expected_events += evaluation_events.copy()
													expected_events += ["on_log", "on_train_end"]
													return expected_events






						def       lowerCamelCase__     (       self      :			Optional[Any]    ):
													__lowerCamelCase					:		Optional[Any]		   =     self.get_trainer()
													__lowerCamelCase					:		Union[str, Any]		   =     DEFAULT_CALLBACKS.copy() + [ProgressCallback]
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )

													# Callbacks passed at init are added to the default callbacks
													__lowerCamelCase					:		Dict		   =     self.get_trainer(callbacks=[MyTestTrainerCallback]    )
													expected_callbacks.append(UpperCAmelCase    )
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )

													# TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
													__lowerCamelCase					:		Union[str, Any]		   =     self.get_trainer(disable_tqdm=UpperCAmelCase    )
													__lowerCamelCase					:		Optional[int]		   =     DEFAULT_CALLBACKS.copy() + [PrinterCallback]
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )






						def       lowerCamelCase__     (       self      :			Optional[int]    ):
													__lowerCamelCase					:		List[str]		   =     DEFAULT_CALLBACKS.copy() + [ProgressCallback]
													__lowerCamelCase					:		int		   =     self.get_trainer()

													# We can add, pop, or remove by class name
													trainer.remove_callback(UpperCAmelCase    )
													expected_callbacks.remove(UpperCAmelCase    )
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )

													__lowerCamelCase					:		Optional[Any]		   =     self.get_trainer()
													__lowerCamelCase					:		Any		   =     trainer.pop_callback(UpperCAmelCase    )
													self.assertEqual(cb.__class__							,					UpperCAmelCase    )
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )

													trainer.add_callback(UpperCAmelCase    )
													expected_callbacks.insert(0							,					UpperCAmelCase    )
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )

													# We can also add, pop, or remove by instance
													__lowerCamelCase					:		Union[str, Any]		   =     self.get_trainer()
													__lowerCamelCase					:		Union[str, Any]		   =     trainer.callback_handler.callbacks[0]
													trainer.remove_callback(UpperCAmelCase    )
													expected_callbacks.remove(UpperCAmelCase    )
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )

													__lowerCamelCase					:		List[Any]		   =     self.get_trainer()
													__lowerCamelCase					:		Tuple		   =     trainer.callback_handler.callbacks[0]
													__lowerCamelCase					:		Optional[Any]		   =     trainer.pop_callback(UpperCAmelCase    )
													self.assertEqual(UpperCAmelCase							,					UpperCAmelCase    )
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )

													trainer.add_callback(UpperCAmelCase    )
													expected_callbacks.insert(0							,					UpperCAmelCase    )
													self.check_callbacks_equality(trainer.callback_handler.callbacks							,					UpperCAmelCase    )







						def       lowerCamelCase__     (       self      :			List[str]    ):
													import warnings

													# XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
													warnings.simplefilter(action="ignore"							,					category=UpperCAmelCase    )

													__lowerCamelCase					:		Dict		   =     self.get_trainer(callbacks=[MyTestTrainerCallback]    )
													trainer.train()
													__lowerCamelCase					:		str		   =     trainer.callback_handler.callbacks[-2].events
													self.assertEqual(UpperCAmelCase							,					self.get_expected_events(UpperCAmelCase    )    )

													# Independent log/save/eval
													__lowerCamelCase					:		Dict		   =     self.get_trainer(callbacks=[MyTestTrainerCallback]							,					logging_steps=5    )
													trainer.train()
													__lowerCamelCase					:		List[str]		   =     trainer.callback_handler.callbacks[-2].events
													self.assertEqual(UpperCAmelCase							,					self.get_expected_events(UpperCAmelCase    )    )

													__lowerCamelCase					:		List[Any]		   =     self.get_trainer(callbacks=[MyTestTrainerCallback]							,					save_steps=5    )
													trainer.train()
													__lowerCamelCase					:		Any		   =     trainer.callback_handler.callbacks[-2].events
													self.assertEqual(UpperCAmelCase							,					self.get_expected_events(UpperCAmelCase    )    )

													__lowerCamelCase					:		Optional[Any]		   =     self.get_trainer(callbacks=[MyTestTrainerCallback]							,					eval_steps=5							,					evaluation_strategy="steps"    )
													trainer.train()
													__lowerCamelCase					:		Tuple		   =     trainer.callback_handler.callbacks[-2].events
													self.assertEqual(UpperCAmelCase							,					self.get_expected_events(UpperCAmelCase    )    )

													__lowerCamelCase					:		int		   =     self.get_trainer(callbacks=[MyTestTrainerCallback]							,					evaluation_strategy="epoch"    )
													trainer.train()
													__lowerCamelCase					:		Dict		   =     trainer.callback_handler.callbacks[-2].events
													self.assertEqual(UpperCAmelCase							,					self.get_expected_events(UpperCAmelCase    )    )

													# A bit of everything
													__lowerCamelCase					:		Optional[Any]		   =     self.get_trainer(
													    callbacks=[MyTestTrainerCallback]							,					logging_steps=3							,					save_steps=10							,					eval_steps=5							,					evaluation_strategy="steps"							,					)
													trainer.train()
													__lowerCamelCase					:		Union[str, Any]		   =     trainer.callback_handler.callbacks[-2].events
													self.assertEqual(UpperCAmelCase							,					self.get_expected_events(UpperCAmelCase    )    )

													# warning should be emitted for duplicated callbacks
													with patch("transformers.trainer_callback.logger.warning"    ) as warn_mock:
																				__lowerCamelCase					:		List[Any]		   =     self.get_trainer(
																				    callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]							,					)
																				assert str(UpperCAmelCase    ) in warn_mock.call_args[0][0]
"""simple docstring"""





import gc
import threading
import time

import psutil
import torch
class  _snake_case     :






						def __init__(       self      :			str    ):
													__lowerCamelCase					:		Optional[Any]		   =     psutil.Process()
													__lowerCamelCase					:		List[Any]		   =     False






						def       lowerCamelCase__     (       self      :			str    ):
													__lowerCamelCase					:		List[Any]		   =     -1

													while True:
																				__lowerCamelCase					:		Union[str, 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       lowerCamelCase__     (       self      :			Any    ):
													__lowerCamelCase					:		Optional[Any]		   =     True
													__lowerCamelCase					:		Union[str, Any]		   =     threading.Thread(target=self.peak_monitor    )
													__lowerCamelCase					:		Optional[Any]		   =     True
													self.thread.start()







						def       lowerCamelCase__     (       self      :			Dict    ):
													__lowerCamelCase					:		str		   =     False
													self.thread.join()
													return self.cpu_memory_peak

__A  =				PeakCPUMemory()


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



							'''simple docstring'''

							__lowerCamelCase					:		Any		   =     {"time": time.time()}

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

							# CPU mem
							__lowerCamelCase					:		Optional[int]		   =     psutil.Process().memory_info().rss
							cpu_peak_tracker.start()

							# GPU mem
							for i in range(torch.cuda.device_count()      ):
														__lowerCamelCase					:		int		   =     torch.cuda.memory_allocated(_lowerCamelCase      )
							torch.cuda.reset_peak_memory_stats()

							return measures


def    lowercase_   (   _lowerCamelCase:	Optional[Any]      ) ->      Tuple:



							'''simple docstring'''

							__lowerCamelCase					:		Tuple		   =     {"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					:		Tuple		   =     (cpu_peak_tracker.stop() - start_measures["cpu"]) / 2**20

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

							return measures







def    lowercase_   (   _lowerCamelCase:	int						,   _lowerCamelCase:	List[Any]      ) ->      Optional[int]:



							'''simple docstring'''

							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(_lowerCamelCase      )]:.2f}MiB"""      )
														__lowerCamelCase					:		List[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"""      )