I'm using VizTracer to debug performance issues in my PyTorch data loading pipeline. Specifically, I'm using a DataLoader with num_workers > 0 to load data in parallel using multiple subprocesses. However, I noticed that VizTracer doesn't trace these worker subprocesses. After some investigation, it seems that the root issue is that PyTorch uses torch.multiprocessing, which wraps and subclasses multiprocessing.Process. Since VizTracer only seems to detect the standard multiprocessing.Process, it doesn't detect or trace these PyTorch worker processes. Of course, if I instantiate VizTracer on every worker and the main process, they are traced, but this is rather cumbersome and will force me to merge the different output files into one. Is there a way to make VizTracer trace these worker processes? If possible, with the least alterations to the source code and using the CLI. Ideally, I would run the following: python -m viztracer myscript.py and no merging or changes to code would be needed, just as if Pytorch used multiprocessing.Process. What I’ve tried or considered: Wrapping data loading logic with manual multiprocessing.Process, ditching the DataLoader, which VizTracer does detect, but this is reinventing the wheel. Changing the code to have a VizTracer instance on every worker, but as I mentioned, this is not as neat as simply using the CLI command. Monkey-patching torch.multiprocessing.Process to point to multiprocessing.Process, though this feels risky, and does not work anyway (maybe PyTorch detects and patches this, idk). Any advice or clean workarounds would be appreciated! Continue reading...
