def test_pytorch_profiler_trainer_ddp(tmpdir, use_output_filename):
"""Ensure that the profiler can be given to the training and default step are properly recorded. """
if use_output_filename:
output_filename = os.path.join(tmpdir, "profiler.txt")
else:
output_filename = None
profiler = PyTorchProfiler(output_filename=output_filename)
model = BoringModel()
trainer = Trainer(
fast_dev_run=True,
profiler=profiler,
accelerator="ddp",
gpus=2,
)
trainer.fit(model)
enabled = use_output_filename or not use_output_filename and profiler.local_rank == 0
if enabled:
assert len(profiler.summary()) > 0
assert set(profiler.profiled_actions.keys()) == {
'training_step_and_backward', 'validation_step'
}
else:
assert profiler.summary() is None
assert set(profiler.profiled_actions.keys()) == set()
# todo (tchaton) add support for all ranks
if use_output_filename and os.getenv("LOCAL_RANK") == "0":
data = Path(profiler.output_fname).read_text()
assert len(data) > 0
def test_pytorch_profiler_trainer_ddp(tmpdir):
"""Ensure that the profiler can be given to the training and default step are properly recorded. """
pytorch_profiler = PyTorchProfiler(dirpath=None, filename="profiler")
model = BoringModel()
trainer = Trainer(
max_epochs=1,
default_root_dir=tmpdir,
limit_train_batches=2,
limit_val_batches=2,
profiler=pytorch_profiler,
accelerator="ddp",
gpus=2,
)
trainer.fit(model)
assert len(pytorch_profiler.summary()) > 0
assert set(pytorch_profiler.profiled_actions) == {
'training_step_and_backward', 'validation_step'
}
files = sorted(f for f in os.listdir(pytorch_profiler.dirpath)
if "fit" in f)
rank = int(os.getenv("LOCAL_RANK", "0"))
expected = f"fit-profiler-{rank}.txt"
assert files[rank] == expected
path = os.path.join(pytorch_profiler.dirpath, expected)
data = Path(path).read_text("utf-8")
assert len(data) > 0
