Profiler patches torch without reverting it

[awaelchli]

🐛 Bug

LitData patches the worker loop in torch here

https://github.com/Lightning-AI/litdata/blob/5d6b8f962a5dce5efc132131fff46f33c7e24fca/src/litdata/streaming/dataloader.py#L454-L456

but this never gets reverted. Subsequent runs with a dataloader will just happily run this worker loop. This hack should be removed, or if it stays it should have a way for it to remain disabled and be idempotent.

This could lead to hard to debug issues on the user's side.
On our side, this currently leads to interactions in our test suite.

[tchaton]

You are right !

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[Borda]

@bhimrazy is this still an issue?

[bhimrazy]

@Borda Currently, we have a check in place to patch only when profiling is enabled:

litData/src/litdata/streaming/dataloader.py

Lines 473 to 478 in 46a4049

	if self._loader._profile_batches and distributed_env.global_rank == 0 and _VIZ_TRACKER_AVAILABLE:
	from torch.utils.data._utils import worker
	worker._worker_loop = _ProfileWorkerLoop(
	self._loader._profile_batches, self._loader._profile_skip_batches, self._loader._profile_dir
	)

That said, it might still be worth revisiting this to find a cleaner or more robust solution.

[Alphaxiaoteng]

I would like to take this bug. I verified that current main still replaces torch.utils.data._utils.worker._worker_loop for profile_batches and leaves the replacement installed after the iterator is created; the profiling test remains skipped for the same interaction. I found no open PR for this scope.

I plan to keep the fix local to src/litdata/streaming/dataloader.py and tests/streaming/test_dataloader.py: save the original loop, restore it in a finally immediately after multiprocessing workers are started (also covering initialization errors), and unskip/extend the regression test to assert that a subsequent loader sees the original Torch loop. I will run the focused profiling/cProfile tests, the full test file, and the repository style checks. This contribution is AI-assisted and will be reviewed and tested before submission.

[tchaton]

Go for it